From 0aa5e1313b49853ac4488ac1466f5eeb458267e1 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Fri, 5 Jun 2026 20:56:42 +0000 Subject: [PATCH 01/52] Complete legacy AQT deprecation and transition to Qwix/FP8 --- .../base_requirements/requirements.txt | 1 - .../cuda12-requirements.txt | 1 - .../tpu-post-train-requirements.txt | 1 - .../tpu-requirements.txt | 1 - .../requirements/requirements.txt | 1 - .../requirements_decoupled_jax_0_7.1.txt | 1 - src/maxtext/configs/base.yml | 7 +- .../disaggregation/llama3_405b_v6e-16-16.yml | 2 - .../interleaved/llama2_70b_v5e-16.yml | 2 - .../interleaved/llama3_405b_v5e-64.yml | 2 - .../interleaved/llama3_70b_v5e-16.yml | 2 - .../configs/tpu/v5e/llama2_70b_v5e-16.yml | 2 - .../configs/tpu/v5e/llama3_405b_v5e-64.yml | 2 - .../configs/tpu/v5e/llama3_70b_v5e-16.yml | 2 - .../v6e/inference/llama4_maverick_v6e-64.yml | 2 - src/maxtext/configs/types.py | 38 +- src/maxtext/inference/kvcache.py | 59 +-- src/maxtext/inference/maxengine/maxengine.py | 65 +-- src/maxtext/layers/attention_mla.py | 2 +- src/maxtext/layers/attention_op.py | 2 +- src/maxtext/layers/attentions.py | 2 +- src/maxtext/layers/decoders.py | 2 +- src/maxtext/layers/engram.py | 2 +- src/maxtext/layers/initializers.py | 5 +- src/maxtext/layers/learn_to_init_layer.py | 2 +- src/maxtext/layers/linears.py | 41 +- src/maxtext/layers/moe.py | 128 +---- src/maxtext/layers/nnx_decoders.py | 2 +- src/maxtext/layers/quantizations.py | 441 +--------------- src/maxtext/models/deepseek.py | 6 +- src/maxtext/models/gemma.py | 2 +- src/maxtext/models/gemma2.py | 2 +- src/maxtext/models/gemma3.py | 2 +- src/maxtext/models/gemma4.py | 2 +- src/maxtext/models/gemma4_small.py | 2 +- src/maxtext/models/gpt3.py | 2 +- src/maxtext/models/gpt_oss.py | 2 +- src/maxtext/models/llama2.py | 2 +- src/maxtext/models/llama4.py | 2 +- src/maxtext/models/mistral.py | 2 +- src/maxtext/models/mixtral.py | 2 +- src/maxtext/models/models.py | 2 +- src/maxtext/models/olmo3.py | 2 +- src/maxtext/models/qwen2.py | 2 +- src/maxtext/models/qwen3.py | 2 +- src/maxtext/models/qwen3_5.py | 2 +- src/maxtext/models/qwen3_custom.py | 2 +- src/maxtext/models/simple_layer.py | 4 +- src/maxtext/utils/layerwise_quantization.py | 470 ------------------ tests/__init__.py | 52 +- tests/unit/quantizations_test.py | 424 ++++++---------- 51 files changed, 313 insertions(+), 1497 deletions(-) delete mode 100644 src/maxtext/utils/layerwise_quantization.py diff --git a/src/dependencies/requirements/base_requirements/requirements.txt b/src/dependencies/requirements/base_requirements/requirements.txt index 5ba8ee5093..919f2665f7 100644 --- a/src/dependencies/requirements/base_requirements/requirements.txt +++ b/src/dependencies/requirements/base_requirements/requirements.txt @@ -1,5 +1,4 @@ absl-py -aqtp array-record chex cloud-accelerator-diagnostics diff --git a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt index 4e52d9125a..96bd058585 100644 --- a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt @@ -10,7 +10,6 @@ annotated-doc>=0.0.4 annotated-types>=0.7.0 antlr4-python3-runtime>=4.9.3 anyio>=4.13.0 -aqtp>=0.9.0 array-record>=0.8.3 astroid>=4.0.4 astunparse>=1.6.3 diff --git a/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt index b0bd183f81..2bae5fd262 100644 --- a/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt @@ -14,7 +14,6 @@ antlr4-python3-runtime>=4.9.3 anyio>=4.13.0 apache-tvm-ffi>=0.1.11 appnope>=0.1.4 ; sys_platform == 'darwin' -aqtp>=0.9.0 array-record>=0.8.3 astor>=0.8.1 astroid>=4.0.4 diff --git a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt index 26ba4fcdda..6e0474b7ff 100644 --- a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt @@ -10,7 +10,6 @@ annotated-doc>=0.0.4 annotated-types>=0.7.0 antlr4-python3-runtime>=4.9.3 anyio>=4.13.0 -aqtp>=0.9.0 array-record>=0.8.3 astroid>=4.0.4 astunparse>=1.6.3 diff --git a/src/dependencies/requirements/requirements.txt b/src/dependencies/requirements/requirements.txt index 05c2be074b..633dfec057 100644 --- a/src/dependencies/requirements/requirements.txt +++ b/src/dependencies/requirements/requirements.txt @@ -1,5 +1,4 @@ absl-py -aqtp array-record cloud-accelerator-diagnostics cloud-tpu-diagnostics diff --git a/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt b/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt index 8f904a3641..e1cec8bef7 100644 --- a/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt +++ b/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt @@ -1,5 +1,4 @@ absl_py>=2.3.1 -aqtp>=0.9.0 chex>=0.1.90 datasets>=4.2.0 etils>=1.13.0 diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index b69565f9c2..bb5619b236 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -112,7 +112,6 @@ dtype: "bfloat16" # used to configure quantization in the transformer layers, defaults to null implying bf16. # possible alternative settings are as follows: # 'int8' for dynamic range quantization using 8-bits -# 'intmp' for mixed precision quantization for inference as described here: src/maxtext/configs/quantization/readme.md # 'fp8' for 8-bit floating-point gemms on nvidia gpus. # 'nanoo_fp8' for 8-bit floating-point gemms on amd mi300/mi325 gpus. # 'fp8_full' for fp8 quantization with static scaling. @@ -123,10 +122,6 @@ constant_bound_config: "" # https://kolonist26-jax-kr.readthedocs.io/en/latest/jax.lax.html#jax.lax.precision matmul_precision: "default" activations_in_float32: false # sets activations to float32 before nonlinearity it true, else dtype -# used to replicate the quantization scale to avoid the inefficient xla fusion for 2d sharding. -replicate_quant_scale: false -# path to file with quantization config for intmp. -quant_cfg_path: "" quantize_kvcache: false # set to true to quantize kv cache values, defaults to false # valid kv_quant_axis values: # - "" is valid only when quantize_kvcache is false @@ -143,7 +138,7 @@ save_quantized_params_path: "" # when left as is, corresponds to training # accepted values are "inference" model_call_mode: "" -use_qwix_quantization: false # [DEPRECATED: AQT will be removed in a future release. It is strongly recommended to set use_qwix_quantization to true] whether to use qwix for quantization. if set to true, the model will be quantized using qwix. +use_qwix_quantization: true # [DEPRECATED: AQT will be removed in a future release. It is strongly recommended to set use_qwix_quantization to true] whether to use qwix for quantization. if set to true, the model will be quantized using qwix. use_manual_quantization: false # a flag if to use manual quantization for batch split. Only used if use_batch_split_schedule is true. # quantization calibration method used for weights and activations. supported methods can be found in https://github.com/google/qwix/blob/dc2a0770351c740e5ab3cce7c0efe9f7beacce9e/qwix/qconfig.py#l70-l80 weight_quantization_calibration_method: "absmax" diff --git a/src/maxtext/configs/inference/multihost/disaggregation/llama3_405b_v6e-16-16.yml b/src/maxtext/configs/inference/multihost/disaggregation/llama3_405b_v6e-16-16.yml index 2b83fadc56..02402e07dd 100644 --- a/src/maxtext/configs/inference/multihost/disaggregation/llama3_405b_v6e-16-16.yml +++ b/src/maxtext/configs/inference/multihost/disaggregation/llama3_405b_v6e-16-16.yml @@ -5,8 +5,6 @@ sharding_strategy: "experimental" attention: 'dot_product' allow_split_physical_axes: true tokenizer_path: "assets/tokenizer_llama3.tiktoken" -# Used to replicate the quantization scale to avoid the inefficient XLA fusion. -replicate_quant_scale: true inference_server: "ExperimentalMaxtextDisaggregatedServer" diff --git a/src/maxtext/configs/inference/multihost/interleaved/llama2_70b_v5e-16.yml b/src/maxtext/configs/inference/multihost/interleaved/llama2_70b_v5e-16.yml index dba4bc03ce..eb366a7ee5 100644 --- a/src/maxtext/configs/inference/multihost/interleaved/llama2_70b_v5e-16.yml +++ b/src/maxtext/configs/inference/multihost/interleaved/llama2_70b_v5e-16.yml @@ -8,8 +8,6 @@ model_name: "llama2-70b" sharding_strategy: "experimental" attention: 'dot_product' allow_split_physical_axes: true -# Used to replicate the quantization scale to avoid the inefficient XLA fusion. -replicate_quant_scale: true logical_axis_rules: [ ['embed', []], diff --git a/src/maxtext/configs/inference/multihost/interleaved/llama3_405b_v5e-64.yml b/src/maxtext/configs/inference/multihost/interleaved/llama3_405b_v5e-64.yml index b71b7990f1..220c496732 100644 --- a/src/maxtext/configs/inference/multihost/interleaved/llama3_405b_v5e-64.yml +++ b/src/maxtext/configs/inference/multihost/interleaved/llama3_405b_v5e-64.yml @@ -10,8 +10,6 @@ sharding_strategy: "experimental" attention: 'dot_product' allow_split_physical_axes: true tokenizer_path: "assets/tokenizer_llama3.tiktoken" -# Used to replicate the quantization scale to avoid the inefficient XLA fusion. -replicate_quant_scale: true logical_axis_rules: [ ['embed', []], diff --git a/src/maxtext/configs/inference/multihost/interleaved/llama3_70b_v5e-16.yml b/src/maxtext/configs/inference/multihost/interleaved/llama3_70b_v5e-16.yml index 525d30e30c..1b1c24dd86 100644 --- a/src/maxtext/configs/inference/multihost/interleaved/llama3_70b_v5e-16.yml +++ b/src/maxtext/configs/inference/multihost/interleaved/llama3_70b_v5e-16.yml @@ -9,8 +9,6 @@ tokenizer_path: "assets/tokenizer_llama3.tiktoken" sharding_strategy: "experimental" attention: 'dot_product' allow_split_physical_axes: true -# Used to replicate the quantization scale to avoid the inefficient XLA fusion. -replicate_quant_scale: true logical_axis_rules: [ ['embed', []], diff --git a/src/maxtext/configs/tpu/v5e/llama2_70b_v5e-16.yml b/src/maxtext/configs/tpu/v5e/llama2_70b_v5e-16.yml index dba4bc03ce..eb366a7ee5 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_70b_v5e-16.yml +++ b/src/maxtext/configs/tpu/v5e/llama2_70b_v5e-16.yml @@ -8,8 +8,6 @@ model_name: "llama2-70b" sharding_strategy: "experimental" attention: 'dot_product' allow_split_physical_axes: true -# Used to replicate the quantization scale to avoid the inefficient XLA fusion. -replicate_quant_scale: true logical_axis_rules: [ ['embed', []], diff --git a/src/maxtext/configs/tpu/v5e/llama3_405b_v5e-64.yml b/src/maxtext/configs/tpu/v5e/llama3_405b_v5e-64.yml index b71b7990f1..220c496732 100644 --- a/src/maxtext/configs/tpu/v5e/llama3_405b_v5e-64.yml +++ b/src/maxtext/configs/tpu/v5e/llama3_405b_v5e-64.yml @@ -10,8 +10,6 @@ sharding_strategy: "experimental" attention: 'dot_product' allow_split_physical_axes: true tokenizer_path: "assets/tokenizer_llama3.tiktoken" -# Used to replicate the quantization scale to avoid the inefficient XLA fusion. -replicate_quant_scale: true logical_axis_rules: [ ['embed', []], diff --git a/src/maxtext/configs/tpu/v5e/llama3_70b_v5e-16.yml b/src/maxtext/configs/tpu/v5e/llama3_70b_v5e-16.yml index 525d30e30c..1b1c24dd86 100644 --- a/src/maxtext/configs/tpu/v5e/llama3_70b_v5e-16.yml +++ b/src/maxtext/configs/tpu/v5e/llama3_70b_v5e-16.yml @@ -9,8 +9,6 @@ tokenizer_path: "assets/tokenizer_llama3.tiktoken" sharding_strategy: "experimental" attention: 'dot_product' allow_split_physical_axes: true -# Used to replicate the quantization scale to avoid the inefficient XLA fusion. -replicate_quant_scale: true logical_axis_rules: [ ['embed', []], diff --git a/src/maxtext/configs/tpu/v6e/inference/llama4_maverick_v6e-64.yml b/src/maxtext/configs/tpu/v6e/inference/llama4_maverick_v6e-64.yml index 68f6c839e8..1e2bbc635a 100644 --- a/src/maxtext/configs/tpu/v6e/inference/llama4_maverick_v6e-64.yml +++ b/src/maxtext/configs/tpu/v6e/inference/llama4_maverick_v6e-64.yml @@ -9,8 +9,6 @@ base_config: "inference/inference_jetstream.yml" sharding_strategy: "experimental" attention: 'dot_product' allow_split_physical_axes: true -# Used to replicate the quantization scale to avoid the inefficient XLA fusion. -replicate_quant_scale: true logical_axis_rules: [ ['embed', []], diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index b908a39053..60b60a67df 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -86,7 +86,6 @@ class QuantizationType(str, Enum): NONE = "" INT4 = "int4" INT8 = "int8" - INTMP = "intmp" FP8_E5M2 = "fp8_e5m2" FP8_E4M3 = "fp8_e4m3" FP8 = "fp8" @@ -339,7 +338,7 @@ class Checkpointing(BaseModel): ) checkpoint_is_quantized: bool = Field( False, - description="Set to True if reading from a saved AQT quantized checkpoint.", + description="Set to True if reading from a saved quantized checkpoint.", ) save_quantized_params_path: PathStr = Field("", description="Path to save params quantized on the fly.") enable_orbax_v1: bool = Field(False, description="Bool flag for enabling Orbax v1.") @@ -427,16 +426,11 @@ class Quantization(BaseModel): QuantizationType.NONE, description="Activates quantization for transformer layers.", ) - replicate_quant_scale: bool = Field( - False, - description="Replicates quantization scale to avoid inefficient XLA fusion.", - ) - quant_cfg_path: PathStr = Field("", description="Path to the configuration file for 'intmp' quantization.") quantize_kvcache: bool = Field(False, description="If True, quantizes the Key-Value cache.") kv_quant_axis: KvQuantAxis = Field(KvQuantAxis.HEADS_AND_DKV, description="Axes to quantize over for the KV cache.") kv_quant_dtype: Literal["int8", "int4"] = Field("int8", description="Data type for KV cache quantization.") quantization_local_shard_count: int = Field(-1, description="Shards the range finding operation for quantization.") - use_qwix_quantization: bool = Field(False, description="Whether to use qwix for quantization.") + use_qwix_quantization: bool = Field(True, description="Whether to use qwix for quantization.") use_manual_quantization: bool = Field( False, description="Whether to use manual quantization for batch split. Only used if use_batch_split_schedule is True.", @@ -2640,12 +2634,19 @@ def get_num_target_devices(): } self.num_slices = max_utils.get_num_slices(raw_keys_for_num_slices) - # Check for AQT deprecation warning + # Enforce that Qwix is required for non-native/non-TE quantization if self.quantization and not self.use_qwix_quantization: - if self.quantization not in ("fp8", "nanoo_fp8") and not self.quantization.startswith("te_"): - logger.warning( - "WARNING: AQT quantization is deprecated and will be removed in a future release. " - "Please migrate to Qwix by setting use_qwix_quantization=True." + is_native_or_te = self.quantization in ( + QuantizationType.FP8, + QuantizationType.NANOO_FP8, + QuantizationType.FP8_NANO_V2, + QuantizationType.FP8_GPU, + ) or self.quantization.startswith("te_") + if not is_native_or_te: + raise ValueError( + f"Quantization type '{self.quantization}' without Qwix (use_qwix_quantization=False) " + f"is unsupported because legacy AQT has been completely removed. " + f"Please migrate to Qwix by setting use_qwix_quantization=True." ) # Default quantization sharding count to number of local devices if not set. @@ -2851,15 +2852,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self.data_sharding[0].remove("stage") self.data_sharding[0].insert(0, "stage") - # Add sharding for FP8 amax history when using pipeline parallelism. - if self.quantization and self.quantization in ( - "fp8", - "nanoo_fp8", - "fp8_gpu", - "te_fp8_delayedscaling", - ): - self.logical_axis_rules.append(["aqt_amax_history", ("stage",)]) - # H. RUN ALL CROSS-FIELD VALIDATIONS if self.load_parameters_path and self.load_full_state_path: raise ValueError("At most one of `load_parameters_path` or `load_full_state_path` should be set.") @@ -3137,7 +3129,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self.use_grpo = False if self.use_batch_split_schedule: - if self.quantization and not self.quantization == "fp8_full": + if self.quantization and self.quantization != "fp8_full": raise ValueError("Batch split quantization only supports `quantization=fp8_full`") if self.opt_type == "muon" and self.decoder_block not in [ diff --git a/src/maxtext/inference/kvcache.py b/src/maxtext/inference/kvcache.py index ba2266060f..e14692de84 100644 --- a/src/maxtext/inference/kvcache.py +++ b/src/maxtext/inference/kvcache.py @@ -22,9 +22,12 @@ from flax import linen as nn from flax import nnx -from aqt.jax.v2 import config as aqt_config -from aqt.jax.v2.aqt_tensor import QTensor as KVTensor -from aqt.jax.v2.flax import aqt_flax + +class KVTensor: + + def __init__(self, *args, **kwargs): + raise NotImplementedError("KV Cache quantization is not supported because AQT is deprecated.") + from maxtext.layers import nnx_wrappers from maxtext.layers.initializers import variable_to_logically_partitioned @@ -96,56 +99,10 @@ def einsum_fn_with_rhs_qtensor( lhs_dequant_mode=None, lhs_calibration_mode=None, ): - """einsum function where QTensor is the right-hand-side""" - # Assumes kv is already quantized. - einsum = jnp.einsum - if self.dtype != jnp.float8_e4m3fn: - num_bits = 4 if self.dtype == jnp.int4 else 8 - kv_cfg = aqt_config.dot_general_make( - lhs_bits=None, - rhs_bits=num_bits, - bwd_bits=None, - use_fwd_quant=False, - ) - else: - kv_cfg = aqt_config.config_fwd_fp8() - - if rhs_dequant_mode: - aqt_config.set_fwd_dequant_mode(kv_cfg, rhs_dequant_mode=rhs_dequant_mode) - if rhs_calibration_mode: - aqt_config.set_fwd_calibration_mode( - kv_cfg, - rhs_calibration_mode=rhs_calibration_mode, - ) - if lhs_dequant_mode: - aqt_config.set_fwd_dequant_mode(kv_cfg, lhs_dequant_mode=lhs_dequant_mode) - if lhs_calibration_mode: - aqt_config.set_fwd_calibration_mode( - kv_cfg, - lhs_calibration_mode=lhs_calibration_mode, - ) - einsum = aqt_flax.AqtEinsum( - rhs_quant_mode=aqt_flax.QuantMode.TRAIN, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.NONE, - cfg=kv_cfg, - ) - return einsum + raise NotImplementedError("KV Cache quantization is not supported because AQT is deprecated.") def einsum_fn_with_rhs_qtensor_and_dequant(self): - """Get einstein summation for different dequant modes.""" - if self.dtype == jnp.float8_e4m3fn: - return self.einsum_fn_with_rhs_qtensor( - lhs_dequant_mode=aqt_config.DequantMode.THIS_INPUT, - lhs_calibration_mode=aqt_config.CalibrationMode.REMAINING_AXIS, - rhs_dequant_mode=aqt_config.DequantMode.OTHER_INPUT, - rhs_calibration_mode=aqt_config.CalibrationMode.REMAINING_AXIS, - ) - else: - return self.einsum_fn_with_rhs_qtensor( - rhs_dequant_mode=aqt_config.DequantMode.OTHER_INPUT, - rhs_calibration_mode=aqt_config.CalibrationMode.REMAINING_AXIS, - ) + raise NotImplementedError("KV Cache quantization is not supported because AQT is deprecated.") def kv_cache_as_linen( diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 080104ee89..74709cd28a 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -325,9 +325,8 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar if self.model.quant and self.config.checkpoint_is_quantized: print("Loading from the quantized checkpoint...") - self.model.quant.quant_mode = quantizations.get_quant_mode("serve") - rng1, rng2, rng3 = jax.random.split(rng, 3) + rng1, rng2 = jax.random.split(rng, 2) if params: print("Resharding given params") init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.config, False, rng) @@ -372,7 +371,10 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar ) if self.model.quant and not self.config.checkpoint_is_quantized: - params = self.quantize_params(state, rng3) + raise ValueError( + "On-the-fly parameter quantization is not supported for the modern backends. " + "Please load a pre-quantized checkpoint." + ) else: params = state.params @@ -540,60 +542,6 @@ def unapply_adapter(self, base_params, adapter_config, adapter_params): else: lora_utils.unapply_lora_from_base_params(base_params, adapter_params, lora_scale_factor) - def quantize_params(self, state, rng: PRNGKeyType | None = None): - """Forward pass to quantize decode params.""" - if rng is None: - rng = jax.random.PRNGKey(0) - if self.config.pure_nnx: - # NNX takes a different code path: convert-on-load lives in `_load_params_nnx` - # via `_convert_and_quantize_nnx`, which runs the dummy forward against a - # CONVERT-mode model and transfers `qrhs.frozen` into the SERVE model. - # The standalone `quantize_params(state, rng)` API expects a Linen-shape - # `state.params` dict and isn't reachable on the NNX pathway in maxengine - # (load_params already dispatched to _load_params_nnx). - raise NotImplementedError( - "Use load_params() on NNX — the convert step runs inside _load_params_nnx via " - "_convert_and_quantize_nnx. quantize_params(state, rng) is the Linen API." - ) - - self.model.quant.quant_mode = quantizations.get_quant_mode("convert") - - @jax.jit - def model_apply(_p, _rng): - image_shape = mm_processor.get_dummy_image_shape_for_init( - model_name=self.config.model_name, - batch_size=self.config.micro_batch_size_to_train_on, - ) - audio_shape = mm_processor.get_dummy_audio_shape_for_init(self.config) - return self.model.apply( - _p | {"aqt": {}}, - jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - jnp.ones((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - encoder_images=jnp.ones(image_shape, dtype=jnp.float32) if self.config.use_multimodal else None, - # encoder_image_masks indicates valid tiles if image tiling + padding is used in vision encoder input. - encoder_image_masks=jnp.ones(image_shape[:2], dtype=jnp.int32) - if self.config.use_multimodal and "llama4" in self.config.model_name - else None, - encoder_audios=jnp.ones(audio_shape, dtype=jnp.float32) if self.config.use_audio else None, - decoder_segment_ids=jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - enable_dropout=False, - model_mode=MODEL_MODE_PREFILL, - rngs={"params": _rng}, - mutable=True, - ) - - _, new_vars = model_apply(state.params, rng) - # Remove param values which have corresponding qtensors in aqt to save memory. - params = {} - params["aqt"] = new_vars["aqt"] - params["params"] = quantizations.remove_quantized_params(state.params["params"], new_vars["aqt"]) - self.abstract_params = jax.tree_util.tree_map( - lambda x: jax.ShapeDtypeStruct(shape=x.shape, dtype=x.dtype, sharding=x.sharding), - params, - ) - maxtext_utils.save_quantized_checkpoint_if_configured(self.config, params) - self.model.quant.quant_mode = quantizations.get_quant_mode("serve") - return params def _maybe_stack_prefill_result_cache(self, cache): """Stack the caches across the layers.""" @@ -2038,8 +1986,7 @@ def set_engine_vars_from_base_engine( """Set internal vars from base_engine, which has already loaded the checkpoint and has sharding, mesh, and kv cache related vars set. """ - if base_engine.model.quant: - engine.model.quant.quant_mode = base_engine.model.quant.quant_mode + engine.state_mesh_annotations = base_engine.state_mesh_annotations engine.abstract_params = base_engine.abstract_params engine.kv_cache_annotations = maxtext_utils.get_kv_cache_annotations(engine.model, engine.config, rng, engine.mesh) # pylint: disable=protected-access diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py index df7fd16ea2..11449e8f72 100644 --- a/src/maxtext/layers/attention_mla.py +++ b/src/maxtext/layers/attention_mla.py @@ -68,7 +68,7 @@ from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned from maxtext.layers.linears import DenseGeneral from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.inference import kvcache from maxtext.inference.kvcache import KVQuant from maxtext.utils.sharding import create_sharding diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index b3c3f296f4..d1d90a12d8 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -70,7 +70,7 @@ from maxtext.kernels.attention.ragged_attention import ragged_mha from maxtext.layers import nnx_wrappers from maxtext.layers.initializers import variable_to_logically_partitioned -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_pspec import numpy as np diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 93c54e25a6..1fe157ccb3 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -64,7 +64,7 @@ from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned, default_bias_init from maxtext.layers.linears import DenseGeneral, canonicalize_tuple, normalize_axes from maxtext.layers.normalizations import RMSNorm, Qwen3NextRMSNorm, GlobalRMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.inference import kvcache from maxtext.inference.kvcache import KVQuant from maxtext.utils.sharding import maybe_shard_with_logical, create_sharding diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 46363dbf70..ac7578d0e1 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -37,7 +37,7 @@ from maxtext.layers.attentions import attention_as_linen from maxtext.layers.embeddings import attend_on_embedding, embed_as_linen, positional_embedding_as_linen from maxtext.layers.normalizations import rms_norm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.models import ( deepseek, deepseek_batchsplit, diff --git a/src/maxtext/layers/engram.py b/src/maxtext/layers/engram.py index 3b2eb4e2b5..855ba2ee79 100644 --- a/src/maxtext/layers/engram.py +++ b/src/maxtext/layers/engram.py @@ -30,7 +30,7 @@ from maxtext.layers.initializers import NdInitializer, nd_dense_init from maxtext.layers.linears import DenseGeneral from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant import numpy as np import sympy import tokenizers diff --git a/src/maxtext/layers/initializers.py b/src/maxtext/layers/initializers.py index bbc6605057..97543eb87d 100644 --- a/src/maxtext/layers/initializers.py +++ b/src/maxtext/layers/initializers.py @@ -20,7 +20,6 @@ from flax import linen as nn from flax import nnx -from aqt.jax.v2 import aqt_tensor from maxtext.common.common_types import Array, DType, Shape, PRNGKey @@ -68,7 +67,7 @@ def variable_to_logically_partitioned(variable: nnx.Variable): present, it wraps the variable's value in `nn.LogicallyPartitioned` to apply the specified sharding constraints. - It handles special cases for `aqt_tensor.QTensor` and variables of type + It handles special cases for variables of type `_overwrite_with_gradient` by returning their values directly without wrapping. @@ -79,8 +78,6 @@ def variable_to_logically_partitioned(variable: nnx.Variable): The variable's value, potentially wrapped in `nn.LogicallyPartitioned`. """ val = variable.get_value() - if isinstance(val, aqt_tensor.QTensor): - return val if variable.type.__name__ == "_overwrite_with_gradient": return val diff --git a/src/maxtext/layers/learn_to_init_layer.py b/src/maxtext/layers/learn_to_init_layer.py index 2530c17336..368f69674e 100644 --- a/src/maxtext/layers/learn_to_init_layer.py +++ b/src/maxtext/layers/learn_to_init_layer.py @@ -24,7 +24,7 @@ from maxtext.common.common_types import DType, ShardMode, Array from maxtext.layers.nnx_wrappers import ToNNX -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.initializers import NdInitializer, nd_dense_init from maxtext.utils import max_logging, max_utils diff --git a/src/maxtext/layers/linears.py b/src/maxtext/layers/linears.py index cc26673c5c..5e6279e06a 100644 --- a/src/maxtext/layers/linears.py +++ b/src/maxtext/layers/linears.py @@ -31,10 +31,10 @@ from maxtext.common.common_types import DecoderBlockType, ShardMode, DType, Array, Config from maxtext.common.common_types import MODEL_MODE_PREFILL -from maxtext.layers import nnx_wrappers, quantizations +from maxtext.layers import nnx_wrappers from maxtext.layers import normalizations from maxtext.layers.initializers import NdInitializer, nd_dense_init, default_bias_init, variable_to_logically_partitioned -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils.sharding import maybe_shard_with_logical @@ -157,17 +157,16 @@ def __init__( kernel_in_axis = np.arange(len(self.axis)) kernel_out_axis = np.arange(len(self.axis), len(self.axis) + len(self.out_features_shape)) - if not quantizations.in_serve_mode(self.quant): - self.kernel = nnx.Param( - self.kernel_init( - rngs.params(), - kernel_shape, - self.weight_dtype, - kernel_in_axis, - kernel_out_axis, - ), - sharding=self.kernel_axes, - ) + self.kernel = nnx.Param( + self.kernel_init( + rngs.params(), + kernel_shape, + self.weight_dtype, + kernel_in_axis, + kernel_out_axis, + ), + sharding=self.kernel_axes, + ) if self.use_bias: bias_axes = self.kernel_axes[-len(self.out_features_shape) :] @@ -216,16 +215,12 @@ def __call__(self, inputs: Array, _initializing: bool = False, out_sharding: Nam f"does not match expected input feature size {self.in_features_shape[i]}" ) - if quantizations.in_serve_mode(self.quant): - kernel_shape = self.in_features_shape + self.out_features_shape - kernel = jnp.zeros(kernel_shape, dtype=self.dtype) - else: - kernel = self.kernel[...] - # Move logit_dense kernel to device if parameter offloading is enabled - if self.parameter_memory_host_offload: - max_logging.log("linear.py: Moving parameter logits_dense kernel to device") - kernel = jax.device_put(kernel, max_utils.device_space()) - kernel = jnp.asarray(kernel, self.dtype) + kernel = self.kernel[...] + # Move logit_dense kernel to device if parameter offloading is enabled + if self.parameter_memory_host_offload: + max_logging.log("linear.py: Moving parameter logits_dense kernel to device") + kernel = jax.device_put(kernel, max_utils.device_space()) + kernel = jnp.asarray(kernel, self.dtype) # out_sharding should be None for auto mesh axis if self.shard_mode != ShardMode.EXPLICIT: diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index abcced3a6a..e222831388 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -13,6 +13,8 @@ # limitations under the License. +# pylint: disable=assignment-from-none + """MoE related Layers.""" import enum @@ -21,7 +23,7 @@ import random from typing import Iterable, Optional, Tuple, Union -from aqt.jax.v2 import aqt_tensor as aqt + from flax import nnx from flax import struct import jax @@ -223,7 +225,7 @@ def __init__( kernel_axes: Tuple[Optional[str], ...] = (), use_bias: bool = False, score_func: str = "", - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, shard_mode: ShardMode = ShardMode.AUTO, matmul_precision: str = "default", ): @@ -266,17 +268,16 @@ def __init__( kernel_in_axis = np.arange(len(self.axis)) kernel_out_axis = np.arange(len(self.axis), len(self.axis) + len(self.out_features_shape)) - if not quantizations.in_serve_mode(self.quant): - self.kernel = nnx.Param( - self.kernel_init( - rngs.params(), - kernel_shape, - self.weight_dtype, - kernel_in_axis, - kernel_out_axis, - ), - out_sharding=self.kernel_axes, - ) + self.kernel = nnx.Param( + self.kernel_init( + rngs.params(), + kernel_shape, + self.weight_dtype, + kernel_in_axis, + kernel_out_axis, + ), + out_sharding=self.kernel_axes, + ) if self.use_bias: bias_axes = self.kernel_axes[-len(self.out_features_shape) :] @@ -309,11 +310,7 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax. inputs = jnp.asarray(inputs, self.dtype) norm_axis = linears.normalize_axes(self.axis, inputs.ndim) - if quantizations.in_serve_mode(self.quant): - kernel_shape = self.in_features_shape + self.out_features_shape - kernel = jnp.zeros(kernel_shape, dtype=self.dtype) - else: - kernel = self.kernel[...] + kernel = self.kernel[...] kernel = jnp.asarray(kernel, self.dtype) contract_ind = tuple(range(0, len(norm_axis))) @@ -360,7 +357,7 @@ def __init__( intermediate_dim: int = 2048, weight_dtype: ctypes.DType = jnp.float32, dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, ): """Initializes the RoutedMoE module. @@ -473,14 +470,7 @@ def __init__( kernel_in_axis = np.arange(1) kernel_out_axis = np.arange(1, 2) - if quantizations.in_serve_mode(self.quant): - # During aqt convert state we delete kernel weight from params to save - # memory. Instead they are retrieved from the tensors stored in the 'aqt' - # collection. - self.wi_0 = jnp.zeros((num_experts, self.moe_expert_input_dim, intermediate_dim)) - self.wi_1 = jnp.zeros((num_experts, self.moe_expert_input_dim, intermediate_dim)) - self.wo = jnp.zeros((num_experts, intermediate_dim, self.moe_expert_input_dim)) - elif self.config.prefuse_moe_weights and self.config.attention == "vllm_rpa": + if self.config.prefuse_moe_weights and self.config.attention == "vllm_rpa": # Pad model dimension in Fused MoE weight kernels for GMM_v2 execution. moe_intermediate_dim = ( self.config.padded_base_moe_mlp_dim @@ -1128,10 +1118,6 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments, min(tiling[2], n), ) rhs_inputs = kernel - if isinstance(kernel, aqt.QTensor): - if kernel.bias or kernel.sparsity_mask or len(kernel.scale) > 1: - raise ValueError("Unsupported usecase for ragged_dot with quantized kernel.") - rhs_inputs = kernel.qvalue if self.config.use_qwix_quantization: # Use full contraction for QWIX quantization to allow quantization # fusion (max reduce over contracting dimension). @@ -1150,16 +1136,6 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments, group_sizes=group_sizes, preferred_element_type=self.dtype, ) - if isinstance(kernel, aqt.QTensor): - # Multiply outputs by the kernely scale - scales = jnp.take(kernel.scale[0].squeeze(), indices=expert_assignments, axis=0) - if padding_amount > 0: - scales = jax.lax.pad( - scales, - jnp.array(0.0, dtype=scales.dtype), - [(0, padding_amount, 0), (0, 0, 0)], - ) - output *= scales return output def get_tokamax_group_sizes(group_sizes, inputs, kernel): @@ -1273,15 +1249,6 @@ def explicitly_weight_ag(shard_exp_on_fsdp): return True return False - def maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec): - if isinstance(w0_kernel, aqt.QTensor): - w0_pspec = aqt.partition_spec(w0_pspec, (1,), w0_kernel.dtype, use_bias=False) - if isinstance(w1_kernel, aqt.QTensor): - w1_pspec = aqt.partition_spec(w1_pspec, (1,), w1_kernel.dtype, use_bias=False) - if isinstance(wo_kernel, aqt.QTensor): - wo_pspec = aqt.partition_spec(wo_pspec, (1,), wo_kernel.dtype, use_bias=False) - return w0_pspec, w1_pspec, wo_pspec - def get_routed_moe_shardings(is_batch_sharded_by_expert): if is_batch_sharded_by_expert: batch_logical_axis = "activation_batch" @@ -1354,7 +1321,6 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert): w1_bias_pspec, wo_bias_pspec, ) = get_routed_moe_shardings(is_batch_sharded_by_expert) - w0_pspec, w1_pspec, wo_pspec = maybe_aqt_partition(w0_kernel, w0_pspec, w1_kernel, w1_pspec, wo_kernel, wo_pspec) def route(x, logits, pre_bias_logits, rngs): """Performs both across device and within device token routing/sorting""" @@ -1950,14 +1916,13 @@ def get_einsum( if self.quant: - def aqt_einsum(*args, **kwargs): # pylint: disable=unused-argument - # simply skip kwargs, since aqt einsum doesn't support any kwargs + def quant_einsum(*args, **kwargs): # pylint: disable=unused-argument + # simply skip kwargs, since einsum doesn't support any kwargs # like precision - is_aqt = not isinstance(self.quant, quantizations.Fp8Quantization) - kw = {"mesh_axes": rhs_mesh_axes} if is_aqt else {"dtype": self.dtype} + kw = {"dtype": self.dtype} return self.quant.einsum(**kw)(*args) # pytype: disable=attribute-error - einsum_op = aqt_einsum + einsum_op = quant_einsum else: einsum_op = jnp.einsum return einsum_op @@ -2334,35 +2299,6 @@ def fused_moe_matmul( output = jnp.reshape(output_2d, (batch_size, seq_len, emb_dim)) return output, None, None - def retrieve_quantized_weight( - self, - inputs, - gate_logits, - pre_bias_logits, - w0_kernel, - w1_kernel, - wo_kernel, - w0_bias, - w1_bias, - wo_bias, - ) -> tuple[aqt.QTensor, aqt.QTensor, aqt.QTensor]: - """Retrieve quantized weights.""" - # This is called only during tracing. This is to invoke creation of - # quantized tensor inside AqtEinsum. After jit, this will become no-op and - # will not affect performance. - _ = self.dense_matmul( - inputs, gate_logits, pre_bias_logits, w0_kernel, w1_kernel, wo_kernel, w0_bias, w1_bias, wo_bias - ) - - w0_kernel = self.variables["aqt"]["AqtEinsum_0"]["AqtDotGeneral_0"]["qrhs"]["frozen"] - w1_kernel = self.variables["aqt"]["AqtEinsum_1"]["AqtDotGeneral_0"]["qrhs"]["frozen"] - wo_kernel = self.variables["aqt"]["AqtEinsum_2"]["AqtDotGeneral_0"]["qrhs"]["frozen"] - - w0_kernel = max_utils.unbox_logicallypartioned(w0_kernel) - w1_kernel = max_utils.unbox_logicallypartioned(w1_kernel) - wo_kernel = max_utils.unbox_logicallypartioned(wo_kernel) - return w0_kernel, w1_kernel, wo_kernel - def __call__( self, inputs: jax.Array, gate_inputs: jax.Array | None = None, out_sharding: NamedSharding | None = None ) -> tuple[jax.Array, Optional[jax.Array], Optional[jax.Array]]: @@ -2403,19 +2339,7 @@ def __call__( output, lb_loss, bias_updates = self.fused_moe_matmul( inputs, gate_logits, wo_kernel, w0_kernel=w0_kernel, w1_kernel=w1_kernel, fused_kernel=fused_kernel ) - elif cfg.sparse_matmul: - if quantizations.in_serve_mode(self.quant): - w0_kernel, w1_kernel, wo_kernel = self.retrieve_quantized_weight( - inputs, - gate_logits, - pre_bias_logits, - w0_kernel, - w1_kernel, - wo_kernel, - w0_bias, - w1_bias, - wo_bias, - ) + output, lb_loss, bias_updates = self.sparse_matmul( inputs, gate_logits, pre_bias_logits, w0_kernel, w1_kernel, wo_kernel, w0_bias, w1_bias, wo_bias ) @@ -2438,7 +2362,7 @@ def __init__( rngs: nnx.Rngs, weight_dtype: ctypes.DType = jnp.float32, dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, ): """Initializes the RoutedAndSharedMoE module. @@ -2525,7 +2449,7 @@ def get_gate_logit( kernel_axes: Tuple[Optional[str], ...] = (), use_bias: bool = False, score_func: str = "", - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, matmul_precision: str = "default", name: Optional[str] = None, ): @@ -2565,7 +2489,7 @@ def get_routed_moe( intermediate_dim: int = 2048, weight_dtype: ctypes.DType = jnp.float32, dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, name: Optional[str] = None, ): """Creates a RoutedMoE Linen module.""" @@ -2596,7 +2520,7 @@ def get_routed_and_shared_moe( kernel_axes: Tuple[Optional[str], ...], weight_dtype: ctypes.DType = jnp.float32, dtype: ctypes.DType = jnp.float32, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, name: Optional[str] = None, ): """Creates a RoutedAndSharedMoE Linen module.""" diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index de48b60830..89dec7042c 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -42,7 +42,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.embeddings import Embed, PositionalEmbedding, attend_on_embedding from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.models import ( deepseek, deepseek_batchsplit, diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 9907deeb97..6917e27127 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -15,17 +15,10 @@ """Quantization library.""" import functools -import json import qwix.pallas as qpl -import re -from typing import Tuple, Sequence, Callable +from typing import Tuple, Callable from dataclasses import dataclass -from aqt.jax.v2 import config as aqt_config -from aqt.jax.v2 import aqt_tensor -from aqt.jax.v2.flax import aqt_flax -from aqt.jax.v2 import tiled_dot_general -from aqt.jax.v2 import calibration import qwix from qwix._src.core import dot_general_qt @@ -33,14 +26,13 @@ import jax import jax.numpy as jnp -from jax.tree_util import tree_flatten_with_path, tree_unflatten from flax.linen import fp8_ops from flax.linen import initializers as flax_initializers import flax.linen as nn from maxtext.common.common_types import DType, Config -from maxtext.inference.kvcache import KVQuant + # Params used to define mixed precision quantization configs DEFAULT = "__default__" # default config @@ -62,141 +54,6 @@ def einsum(self, dtype: DType = jnp.float32): """Placeholder for einsum implementation in subclasses.""" -def _tiling_fn(lhs, rhs, dimension_numbers, tile_size): - """apply tiling function""" - del lhs, rhs - - (lhs_ca, rhs_ca), _ = dimension_numbers - ret = tiled_dot_general.Cfg( - lhs=tiled_dot_general.TensorTiling(contraction_axes=[], remaining_axes=[]), - rhs=tiled_dot_general.TensorTiling(contraction_axes=[], remaining_axes=[]), - ) - - for lhs_idx, rhs_idx in zip(lhs_ca, rhs_ca): - ret.lhs.contraction_axes.append(tiled_dot_general.AxisTiling(axis=lhs_idx, tile_size=tile_size, tile_count=None)) - ret.rhs.contraction_axes.append(tiled_dot_general.AxisTiling(axis=rhs_idx, tile_size=tile_size, tile_count=None)) - - return ret - - -def _rhs_axis_metadata_wrapper( - x: jnp.ndarray, - tile_map, - no_sharding_axis: Sequence[int], - mesh_axes: Tuple[str, ...], - is_tiled: bool, - replicate_scale: bool = False, -): - """right-hand-side axis metadata wrapper""" - if replicate_scale: - # Temporarily using the shape to identify the scale. - # TODO: remove the replication once the 2d sharding quantization - # works as expected. - if len(x.shape) == 1: - return nn.with_logical_partitioning((lambda: x), tuple(None for _ in mesh_axes))() - - mesh_axes = list(mesh_axes) - if is_tiled: - # tile_map is a mapping between original rank and a list of new, tiled rank. - if len(mesh_axes) < len(tile_map): - mesh_axes = [None] * (len(tile_map) - len(mesh_axes)) + mesh_axes - new_mesh_axes = [None] * len(x.shape) - for orig_rank, new_rank in tile_map.items(): - assert new_rank - assert len(new_rank) <= 2 - new_mesh_axes[new_rank[-1]] = mesh_axes[orig_rank] - mesh_axes = new_mesh_axes - - if mesh_axes is not None and len(mesh_axes) > 0: - for no_shard_idx in no_sharding_axis: - if no_shard_idx < len(mesh_axes): - mesh_axes[no_shard_idx] = None - - return nn.with_logical_partitioning((lambda: x), mesh_axes)() - - -@dataclass -class AqtQuantization: - """Configures AQT quantization github.com/google/aqt.""" - - quant_dg: aqt_config.DotGeneral - quant_mode: aqt_flax.QuantMode = aqt_flax.QuantMode.TRAIN - replicate_scale: bool = False - - def _get_mixed_precision_cfg(self): - """get configuration for mixed precision""" - quant_dg = None - is_tiled = False - tiling_fn = None - # pylint: disable=protected-access - module_path = "/".join(nn.module._context.module_stack[-1].path) - tile_size = -1 - for layer_name_re, layer_quant_dg in self.quant_dg.items(): - if re.fullmatch(layer_name_re, module_path): - quant_dg, tile_size = layer_quant_dg - if quant_dg is None: - quant_dg, tile_size = self.quant_dg[DEFAULT] - if tile_size != -1: - is_tiled = True - tiling_fn = functools.partial(_tiling_fn, tile_size=tile_size) - return quant_dg, is_tiled, tiling_fn - - def _get_rhs_axis_metadata_wrapper( - self, mesh_axes: Tuple[str, ...] = (), is_tiled: bool = False, replicate_scale: bool = False - ): - if self.quant_mode == aqt_flax.QuantMode.CONVERT: - return None - return functools.partial( - _rhs_axis_metadata_wrapper, mesh_axes=mesh_axes, is_tiled=is_tiled, replicate_scale=replicate_scale - ) - - def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): - """Returns dot_general configured with aqt params.""" - if isinstance(self.quant_dg, dict): - quant_dg, is_tiled, tiling_fn = self._get_mixed_precision_cfg() - else: - quant_dg, is_tiled, tiling_fn = self.quant_dg, False, None - rhs_axis_metadata_wrapper = self._get_rhs_axis_metadata_wrapper( - mesh_axes, is_tiled, replicate_scale=self.replicate_scale - ) - # module_path = "/".join(nn.module._context.module_stack[-1].path) - # print(f"quant_dg: {quant_dg}, is_tiled: {is_tiled}, module_path: {module_path}") - aqt_dg_cls = functools.partial( - aqt_flax.AqtDotGeneral, - quant_dg, - rhs_quant_mode=self.quant_mode, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.CALIBRATION_AND_VALUE, - rhs_axis_metadata_wrapper=rhs_axis_metadata_wrapper, - use_legacy_freezer=False, - tiling_fn=tiling_fn, - ) - return aqt_dg_cls - - def einsum(self, mesh_axes: Tuple[str, ...] = ()): - """Returns einsum configured with aqt params.""" - if isinstance(self.quant_dg, dict): - quant_dg, is_tiled, tiling_fn = self._get_mixed_precision_cfg() - else: - quant_dg, is_tiled, tiling_fn = self.quant_dg, False, None - - rhs_axis_metadata_wrapper = self._get_rhs_axis_metadata_wrapper( - mesh_axes, is_tiled, replicate_scale=self.replicate_scale - ) - aqt_einsum = functools.partial( - aqt_flax.AqtEinsum( - cfg=quant_dg, - rhs_quant_mode=self.quant_mode, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.CALIBRATION_AND_VALUE, - rhs_axis_metadata_wrapper=rhs_axis_metadata_wrapper, - use_legacy_freezer=False, - tiling_fn=tiling_fn, - ) - ) - return aqt_einsum - - @dataclass class QwixQuantization: """Configures Qwix quantization github.com/google/qwix, for training only.""" @@ -379,256 +236,22 @@ def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): return nn.NANOOFp8DotGeneralOp -def _get_int8_quant_config(config): - drhs_bits = None - drhs_accumulator_dtype = None - drhs_local_aqt = None - if config.quantization_local_shard_count != 0: - drhs_bits = 8 - drhs_accumulator_dtype = jnp.int32 - drhs_local_aqt = aqt_config.LocalAqt(contraction_axis_shard_count=config.quantization_local_shard_count) - return aqt_config.config_v3( - fwd_bits=8, - dlhs_bits=8, - drhs_bits=drhs_bits, - rng_type="jax.uniform", - dlhs_local_aqt=None, - drhs_local_aqt=drhs_local_aqt, - fwd_accumulator_dtype=jnp.int32, - dlhs_accumulator_dtype=jnp.int32, - drhs_accumulator_dtype=drhs_accumulator_dtype, - ) - - -@dataclass(frozen=True) -class ConstantBoundConfig: - fwd_lhs_bound: float | None = None - fwd_rhs_bound: float | None = None - dlhs_lhs_bound: float | None = None - dlhs_rhs_bound: float | None = None - drhs_lhs_bound: float | None = None - drhs_rhs_bound: float | None = None - - -def _build_const_scale_config( - aqt_dg: aqt_config.DotGeneral, - cst_bound_config: ConstantBoundConfig, -) -> aqt_config.DotGeneral: - """Build a constant scale config for AQT dot general. - - Args: - aqt_dg: The AQT dot general config. - cst_bound_config: The constant bound config. - - Returns: - The AQT dot general config with constant scale config. - """ - if cst_bound_config.fwd_lhs_bound is not None: - aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.fwd_lhs_bound - ) - if cst_bound_config.fwd_rhs_bound is not None: - aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.fwd_rhs_bound - ) - if cst_bound_config.dlhs_lhs_bound: - aqt_dg.dlhs.dg_quantizer.lhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.dlhs_lhs_bound - ) - - if cst_bound_config.dlhs_rhs_bound is not None: - aqt_dg.dlhs.dg_quantizer.rhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.dlhs_rhs_bound - ) - - if cst_bound_config.drhs_lhs_bound is not None: - aqt_dg.drhs.dg_quantizer.lhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.drhs_lhs_bound - ) - - if cst_bound_config.drhs_rhs_bound is not None: - aqt_dg.drhs.dg_quantizer.rhs.calibration = functools.partial( - calibration.ConstantCalibration, bound=cst_bound_config.drhs_rhs_bound - ) - - return aqt_dg - - -@dataclass -class PerTensorScales: - fwd_lhs: bool = False - fwd_rhs: bool = False - dlhs_lhs: bool = False - dlhs_rhs: bool = False - drhs_lhs: bool = False - drhs_rhs: bool = False - - -def _build_per_tensor_config( - aqt_dg: aqt_config.DotGeneral, - per_tensor_scales: PerTensorScales, -) -> aqt_config.DotGeneral: - """Build a per tensor config for AQT dot general. - - Args: - aqt_dg: The AQT dot general config. - per_tensor_scales: The per tensor scales config. - - Returns: - The AQT dot general config with per tensor config. - """ - if per_tensor_scales.fwd_lhs: - aqt_dg.fwd.dg_quantizer.lhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.fwd_rhs: - aqt_dg.fwd.dg_quantizer.rhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.dlhs_lhs: - aqt_dg.dlhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.dlhs_rhs: - aqt_dg.dlhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.drhs_lhs: - aqt_dg.drhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" - if per_tensor_scales.drhs_rhs: - aqt_dg.drhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" - return aqt_dg - - -# fp8 training recipe of dynamic scaling with configurable constant_bound_config for static scaling option -def _get_aqt_fp8_default_config(config): - """Get aqt for 8-bit floating point quantization configuration.""" - aqt_dg = aqt_config.config_v4( - fwd_bits="e4m3", - dlhs_bits="e5m2", - drhs_bits="e5m2", - use_dummy_static_bound=False, - fwd_accumulator_dtype=jnp.bfloat16, - dlhs_accumulator_dtype=jnp.bfloat16, - drhs_accumulator_dtype=jnp.bfloat16, - dlhs_use_fwd_quant=False, - drhs_use_fwd_quant=False, - ) - constant_bound_config = None - - if len(config.constant_bound_config) == 6: - fwd_lhs_bound, fwd_rhs_bound, dlhs_lhs_bound, dlhs_rhs_bound, drhs_lhs_bound, drhs_rhs_bound = ( - config.constant_bound_config - ) - constant_bound_config = ConstantBoundConfig( - fwd_lhs_bound=fwd_lhs_bound, - fwd_rhs_bound=fwd_rhs_bound, - dlhs_lhs_bound=dlhs_lhs_bound, - dlhs_rhs_bound=dlhs_rhs_bound, - drhs_lhs_bound=drhs_lhs_bound, - drhs_rhs_bound=drhs_rhs_bound, - ) - aqt_dg = _build_const_scale_config(aqt_dg, constant_bound_config) - - aqt_config.set_stochastic_rounding( - aqt_dg, - vjp_lhs_stochastic_rounding=False, - vjp_rhs_stochastic_rounding=False, - implementation="jax.uniform", - ) - - per_tensor_scales = PerTensorScales( - fwd_lhs=True, - fwd_rhs=True, - dlhs_lhs=True, - dlhs_rhs=True, - drhs_lhs=True, - drhs_rhs=True, - ) - return _build_per_tensor_config(aqt_dg, per_tensor_scales) - - -def _get_aqt_fp8_quant_config(config): - """get aqt for 8-bit floating point quantization configuration""" - cfg = aqt_config.config_v4(fwd_bits="e4m3", dlhs_bits=None, drhs_bits=None, fwd_accumulator_dtype=jnp.bfloat16) - return cfg - - -def _dot_general_make(quant_cfg): - """Create quantization configs for input matrices to a matmul""" - lhs_bits = quant_cfg[_A_BITS] - lhs_scale = quant_cfg[_A_SCALE] - rhs_bits = quant_cfg[_W_BITS] - rhs_scale = quant_cfg[_W_SCALE] - aqt_dg = aqt_config.dot_general_make(lhs_bits=lhs_bits, rhs_bits=rhs_bits) - if lhs_scale < 1.0: - aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=lhs_scale) - if rhs_scale < 1.0: - aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=rhs_scale) - return aqt_dg - - -def _get_default_mp_config(default=None): - default_config = {_W_BITS: None, _A_BITS: None, _W_SCALE: 1.0, _A_SCALE: 1.0, _TILE_SIZE: -1} - if default: - default_config.update(default) - return default_config - - -def _get_mixed_precision_quant_config(mixed_precision_config): - """Set quantization params based on user configuration.""" - ret_config = {} - default_mp_config = _get_default_mp_config(default=mixed_precision_config.get(DEFAULT, None)) - for layer_name_re, layer_quantization_config in mixed_precision_config.items(): - # Make a copy of default_mp_config to avoid updating original dict - quant_config = default_mp_config.copy() - # print(f"Mixed precision config: processing - # {layer_name_re} - {layer_quantization_config}, default config - {quant_config}") - if layer_name_re != DEFAULT: - for k in quant_config: - quant_config[k] = layer_quantization_config.get(k, default_mp_config[k]) - ret_config[layer_name_re] = [_dot_general_make(quant_config), quant_config["tile_size"]] - return ret_config - - def _get_quant_config(config): """Set quantization params based on user configuration.""" if not config.quantization or config.quantization == "": return None - if config.quantization == "int8": - return _get_int8_quant_config(config) - if config.quantization == "intmp": - assert config.quant_cfg_path, "Must specify quant_cfg for mixed precision quantization" - with open(config.quant_cfg_path, "rt", encoding="utf8") as config_file: - mixed_precision_config = json.load(config_file) - return _get_mixed_precision_quant_config(mixed_precision_config) if config.quantization == "fp8": return "fp8" if config.quantization == "nanoo_fp8": return "nanoo_fp8" - if config.quantization == "aqt_fp8": - return _get_aqt_fp8_quant_config(config) - if config.quantization == "aqt_fp8_full": - return _get_aqt_fp8_default_config(config) if config.quantization.startswith("te_"): return config.quantization - - raise ValueError(f"Invalid value configured for quantization {config.quantization}.") - - -def in_convert_mode(quant): - return quant and (quant.quant_mode == aqt_flax.QuantMode.CONVERT) - - -def in_serve_mode(quant): - return quant and (quant.quant_mode == aqt_flax.QuantMode.SERVE) - - -def get_quant_mode(quant_mode_str: str = "train"): - """Set quant mode.""" - if quant_mode_str == "train": - return aqt_flax.QuantMode.TRAIN - elif quant_mode_str == "serve": - return aqt_flax.QuantMode.SERVE - elif quant_mode_str == "convert": - return aqt_flax.QuantMode.CONVERT - raise ValueError(f"Invalid quantization mode {quant_mode_str}.") + return None def configure_quantization(config: Config, quant_mode_str: str = "train"): """Configure quantization based on user config and quant mode.""" + del quant_mode_str # Unused since AQT is removed if config.use_batch_split_schedule and config.quantization: # The older version of batch-split that fully uses qwix quantization. if config.quantization == "fp8_full" and not config.use_manual_quantization: @@ -650,61 +273,15 @@ def configure_quantization(config: Config, quant_mode_str: str = "train"): return NANOOFp8Quantization() elif isinstance(quant_cfg, str) and quant_cfg.startswith("te_"): return TransformerEngineQuantization(config) - quant_mode = get_quant_mode(quant_mode_str) - replicate_scale = config.replicate_quant_scale if config.replicate_quant_scale else False - return AqtQuantization(quant_dg=quant_cfg, quant_mode=quant_mode, replicate_scale=replicate_scale) return None -def match_aqt_and_unquantized_param(aqt_params, params): - """match aqt and unquantized params""" - aqt_param_flat, aqt_tree_def = jax.tree_util.tree_flatten_with_path( - aqt_params, is_leaf=lambda x: isinstance(x, aqt_tensor.QTensor) - ) - param_tree_flat, _ = jax.tree_util.tree_flatten_with_path(params) - aqt_paths = [] - # Original path of quantized AQT param path. - param_paths = [] - - for aqt_k, _ in aqt_param_flat: - index = None - for index, (k, _) in enumerate(param_tree_flat): - path_depth = len(k) - # every quantized parameter has AQT.. as the leaf node - # AqtDotGeneral and AqtEinsum replace leaf node. - # Therefore, leaf node should be ignored for path matching - # Note: Aqt only operates on kernels so don't pop bias parameters. - # Ref: https://github.com/AI-Hypercomputer/maxtext/compare/main...quantize_r1 - if k[: path_depth - 1] == aqt_k[: path_depth - 1] and k[-1].key != "bias": - aqt_paths.append(aqt_k) - param_paths.append(k) - break - assert index is not None - # since the parameter is already added, we can delete it. - param_tree_flat.pop(index) - return jax.tree_util.tree_unflatten(aqt_tree_def, param_paths) - - -def _get_aqt_key_paths(aqt_vars, params): - """Generate a list of paths which have aqt state""" - aqt_to_unquantized_key_path = match_aqt_and_unquantized_param(aqt_vars, params) - aqt_key_paths, _ = jax.tree_util.tree_flatten(aqt_to_unquantized_key_path, is_leaf=lambda x: isinstance(x, tuple)) - return list(aqt_key_paths) - - -def remove_quantized_params(params, aqt_vars): - """Remove param values with aqt tensors to Null to optimize memory.""" - quantized_param_paths = _get_aqt_key_paths(aqt_vars, params) - tree_flat, tree_struct = tree_flatten_with_path(params) - for i, (k, v) in enumerate(tree_flat): - if k in quantized_param_paths: - v = {} - tree_flat[i] = v - return tree_unflatten(tree_struct, tree_flat) - - def configure_kv_quant(config): - return None if not config.quantize_kvcache else KVQuant(config) + if config.quantize_kvcache: + raise ValueError( + "KV cache quantization (quantize_kvcache=True) is no longer supported " + "because Accurate Quantized Training (AQT) has been deprecated and removed from MaxText." + ) class NvidaFp8Provider(qwix.QtProvider): diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index 27e1a6f7ad..3df1182a94 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -63,7 +63,7 @@ def __init__( model_mode: str, mesh: Mesh, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, layer_idx: int = -1, ) -> None: self.config = config @@ -315,7 +315,7 @@ def __init__( model_mode: str, mesh: Mesh, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, layer_idx: int = -1, ) -> None: super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) @@ -404,7 +404,7 @@ def __init__( model_mode: str, mesh: Mesh, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, layer_idx: int = -1, ) -> None: super().__init__(config, model_mode, mesh, rngs, quant, layer_idx) diff --git a/src/maxtext/models/gemma.py b/src/maxtext/models/gemma.py index 84f4f6817d..3f9ff0682e 100644 --- a/src/maxtext/models/gemma.py +++ b/src/maxtext/models/gemma.py @@ -29,7 +29,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/gemma2.py b/src/maxtext/models/gemma2.py index a7315763eb..e3ff18fc32 100644 --- a/src/maxtext/models/gemma2.py +++ b/src/maxtext/models/gemma2.py @@ -30,7 +30,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/gemma3.py b/src/maxtext/models/gemma3.py index 92c46b96f7..4ad017d42e 100644 --- a/src/maxtext/models/gemma3.py +++ b/src/maxtext/models/gemma3.py @@ -29,7 +29,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock, Dropout from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.initializers import variable_to_logically_partitioned from maxtext.utils import max_utils diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py index 626d2ff54c..3ac037a983 100644 --- a/src/maxtext/models/gemma4.py +++ b/src/maxtext/models/gemma4.py @@ -33,7 +33,7 @@ import jax.sharding from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/gemma4_small.py b/src/maxtext/models/gemma4_small.py index ca33470bf2..1ee945420e 100644 --- a/src/maxtext/models/gemma4_small.py +++ b/src/maxtext/models/gemma4_small.py @@ -29,7 +29,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/gpt3.py b/src/maxtext/models/gpt3.py index a6b08d8b24..df16ad1ecc 100644 --- a/src/maxtext/models/gpt3.py +++ b/src/maxtext/models/gpt3.py @@ -35,7 +35,7 @@ from maxtext.layers import linears from maxtext.layers.attentions import AttentionOp, KVQuant from maxtext.layers.initializers import Initializer, NdInitializer, nd_dense_init -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_logging from maxtext.utils import max_utils diff --git a/src/maxtext/models/gpt_oss.py b/src/maxtext/models/gpt_oss.py index e854a75556..7b03cba502 100644 --- a/src/maxtext/models/gpt_oss.py +++ b/src/maxtext/models/gpt_oss.py @@ -35,7 +35,7 @@ from maxtext.layers import quantizations from maxtext.layers.attentions import Attention from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils # ----------------------------------------- diff --git a/src/maxtext/models/llama2.py b/src/maxtext/models/llama2.py index 0c3e0cca7c..260934acc1 100644 --- a/src/maxtext/models/llama2.py +++ b/src/maxtext/models/llama2.py @@ -30,7 +30,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils from maxtext.utils.sharding import create_sharding, maybe_shard_with_logical from maxtext.layers.learn_to_init_layer import apply_lti_modification diff --git a/src/maxtext/models/llama4.py b/src/maxtext/models/llama4.py index 570a266dc7..df81e09b53 100644 --- a/src/maxtext/models/llama4.py +++ b/src/maxtext/models/llama4.py @@ -35,7 +35,7 @@ from maxtext.layers.linears import MlpBlock from maxtext.layers.moe import RoutedAndSharedMoE from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils #### Multi modal model implementation diff --git a/src/maxtext/models/mistral.py b/src/maxtext/models/mistral.py index 49a74c95db..8c4747e166 100644 --- a/src/maxtext/models/mistral.py +++ b/src/maxtext/models/mistral.py @@ -28,7 +28,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout, MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils # ----------------------------------------- diff --git a/src/maxtext/models/mixtral.py b/src/maxtext/models/mixtral.py index faf69273c6..38a9d5df41 100644 --- a/src/maxtext/models/mixtral.py +++ b/src/maxtext/models/mixtral.py @@ -29,7 +29,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import Dropout from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils # ----------------------------------------- diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index 2af0d560da..20a171a8bb 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -33,7 +33,7 @@ from maxtext.layers.embeddings import Embed, embed_as_linen from maxtext.layers.encoders import AudioEncoder, VisionEncoder, audio_encoder_as_linen, vision_encoder_as_linen from maxtext.layers.multi_token_prediction import MultiTokenPredictionBlock, multi_token_prediction_block_as_linen -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.multimodal import processor as mm_processor from maxtext.utils import max_utils diff --git a/src/maxtext/models/olmo3.py b/src/maxtext/models/olmo3.py index fe8a4e489e..3cebd6170f 100644 --- a/src/maxtext/models/olmo3.py +++ b/src/maxtext/models/olmo3.py @@ -36,7 +36,7 @@ from maxtext.layers.attentions import Attention from maxtext.layers.linears import MlpBlock from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils diff --git a/src/maxtext/models/qwen2.py b/src/maxtext/models/qwen2.py index 69555c176d..995e7d4c54 100644 --- a/src/maxtext/models/qwen2.py +++ b/src/maxtext/models/qwen2.py @@ -30,7 +30,7 @@ from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.attentions import Attention from maxtext.layers.linears import MlpBlock from maxtext.utils import max_utils diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index e527b59b24..a255351253 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -37,7 +37,7 @@ from maxtext.layers import quantizations from maxtext.layers.embeddings import Qwen3OmniMoeVisionPosEmbedInterpolate, PositionalEmbedding from maxtext.layers.normalizations import RMSNorm, l2norm, Qwen3NextRMSNorm, Qwen3NextRMSNormGated -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral, MlpBlock from maxtext.layers.moe import RoutedMoE diff --git a/src/maxtext/models/qwen3_5.py b/src/maxtext/models/qwen3_5.py index 759fd180cd..b10a19e84c 100644 --- a/src/maxtext/models/qwen3_5.py +++ b/src/maxtext/models/qwen3_5.py @@ -28,7 +28,7 @@ from maxtext.layers import initializers as max_initializers from maxtext.layers import nnx_wrappers from maxtext.layers.normalizations import Qwen3NextRMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.utils import max_utils from maxtext.models.qwen3 import ( diff --git a/src/maxtext/models/qwen3_custom.py b/src/maxtext/models/qwen3_custom.py index e0ca4bb512..a78beb6b44 100644 --- a/src/maxtext/models/qwen3_custom.py +++ b/src/maxtext/models/qwen3_custom.py @@ -25,7 +25,7 @@ from maxtext.layers import moe from maxtext.layers import nnx_wrappers from maxtext.layers import quantizations -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.layers.attentions import Attention from maxtext.layers.linears import DenseGeneral from maxtext.utils import max_utils diff --git a/src/maxtext/models/simple_layer.py b/src/maxtext/models/simple_layer.py index ac4eb915a8..1e72b38143 100644 --- a/src/maxtext/models/simple_layer.py +++ b/src/maxtext/models/simple_layer.py @@ -37,7 +37,7 @@ def __init__( mesh: Mesh, model_mode: str, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, ) -> None: self.config = config @@ -93,7 +93,7 @@ def __init__( mesh: Mesh, model_mode: str, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, ) -> None: self.config = config diff --git a/src/maxtext/utils/layerwise_quantization.py b/src/maxtext/utils/layerwise_quantization.py deleted file mode 100644 index 3279b59a6f..0000000000 --- a/src/maxtext/utils/layerwise_quantization.py +++ /dev/null @@ -1,470 +0,0 @@ -# Copyright 2023–2026 Google LLC -# -# 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 -# -# https://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. - - -r"""Layerwise quantization for large models - -Provides a utility to load and quantize a checkpoint layer by layer. Currently, it supports DeepSeek-family models only. - -Example cmd: - -python3 -m maxtext.utils.layerwise_quantization src/maxtext/configs/base.yml \ - tokenizer_path=${TOKENIZER_PATH?} load_parameters_path=${LOAD_PARAMS_PATH?} \ - model_name=deepseek2-16b ici_fsdp_parallelism=1 ici_autoregressive_parallelism=1 \ - ici_tensor_parallelism=-1 scan_layers=false weight_dtype=bfloat16 per_device_batch_size=1 \ - attention=dot_product quantization=int8 async_checkpointing=false enable_single_controller=true \ - tokenizer_type=huggingface megablox=false sparse_matmul=false \ - save_quantized_params_path=${SAVE_PARAMS_PATH?} checkpoint_storage_use_ocdbt=False \ - checkpoint_storage_use_zarr3=False - -""" - -import functools -import os -from typing import Any, Sequence - -from absl import app -from aqt.jax.v2 import aqt_tensor -from flax import nnx -from flax.linen import partitioning as nn_partitioning -import jax -import jax.numpy as jnp -from maxtext.common import common_types -from maxtext.common import checkpointing -from maxtext.layers import quantizations -from maxtext.models import deepseek, models -from maxtext.utils import max_logging -from maxtext.utils import max_utils -from maxtext.utils import maxtext_utils -from maxtext.utils import maxtext_utils_nnx -from maxtext.utils import model_creation_utils -import orbax.checkpoint as ocp -from tqdm import tqdm -from maxtext.configs import pyconfig - -IGNORE = ocp.PLACEHOLDER -PRNGKeyType = Any -DictKey = jax.tree_util.DictKey - - -def get_original_path_key(aqt_k_tuple: tuple[DictKey, ...]) -> tuple[DictKey, ...] | None: - """ - Maps an AQT PyTree path (tuple of keys) to its corresponding original parameter path. - Only returns a path if it corresponds to a parameter to be removed. - """ - aqt_k = list(aqt_k_tuple) - str_path = jax.tree_util.keystr(aqt_k_tuple) - if "AqtEinsum_" in str_path: - return None - if "AqtDotGeneral_" not in str_path: - return None - aqt_module_index = -1 - for i, key in enumerate(aqt_k): - if isinstance(key, DictKey) and key.key.startswith("AqtDotGeneral_"): - aqt_module_index = i - break - if aqt_module_index == -1: - return None - if ( - len(aqt_k) > aqt_module_index + 2 - and isinstance(aqt_k[aqt_module_index + 1], DictKey) - and aqt_k[aqt_module_index + 1].key == "qrhs" - and isinstance(aqt_k[aqt_module_index + 2], DictKey) - and aqt_k[aqt_module_index + 2].key == "frozen" - ): - parent_path = tuple(aqt_k[:aqt_module_index]) - return parent_path + (DictKey("kernel"),) - return None - - -def get_quantized_param_paths(aqt_params: Any, params: Any) -> set[tuple[DictKey, ...]]: - """ - Identifies the set of paths in the original params tree that have been quantized. - """ - - def is_qtensor(x): - return isinstance(x, aqt_tensor.QTensor) - - aqt_param_flat, _ = jax.tree_util.tree_flatten_with_path(aqt_params, is_leaf=is_qtensor) - if not aqt_param_flat: - return set() - param_tree_flat_with_path, _ = jax.tree_util.tree_flatten_with_path(params) - params_path_set: set[tuple[DictKey, ...]] = {tuple(k) for k, _ in param_tree_flat_with_path} - original_param_paths_to_remove: set[tuple[DictKey, ...]] = set() - for aqt_k_tuple, _ in aqt_param_flat: - original_k_tuple = get_original_path_key(aqt_k_tuple) - if original_k_tuple is None: - continue - if original_k_tuple in params_path_set: - original_param_paths_to_remove.add(original_k_tuple) - continue - params_keys_str = {jax.tree_util.keystr(k) for k in params_path_set} - raise ValueError( - f"Mapped AQT path {jax.tree_util.keystr(aqt_k_tuple)} to {jax.tree_util.keystr(original_k_tuple)}," - f" but not found in params. Available: {params_keys_str}" - ) - return original_param_paths_to_remove - - -def remove_quantized_params(params: Any, aqt_vars: Any) -> Any: - """Replaces the values in the original params tree that are now quantized with empty dicts.""" - quantized_param_path_set = get_quantized_param_paths(aqt_vars, params) - if not quantized_param_path_set: - return params - - def _map_fn(path, value): - return {} if tuple(path) in quantized_param_path_set else value - - return jax.tree_util.tree_map_with_path(_map_fn, params) - - -# --- Function to restructure NNX-Run AQT tree to match PURE LINEN saved format --- -def insert_deepseekmoeblock_scope(aqt_layer_tree: dict[str, Any]) -> dict[str, Any]: - """ - Moves top-level AqtEinsum_* entries into the existing 'DeepSeekMoeBlock_0' - dict to match the pure Linen AQT structure. - """ - if not isinstance(aqt_layer_tree, dict): - return aqt_layer_tree - - new_tree = dict(aqt_layer_tree) # Start with a copy - - einsum_items = {key: new_tree.pop(key) for key in list(new_tree.keys()) if key.startswith("AqtEinsum_")} - - if einsum_items: - if "DeepSeekMoeBlock_0" not in new_tree: - # This case indicates the MoE block itself was missing, which is unexpected - max_logging.log("Error: 'DeepSeekMoeBlock_0' not found in AQT vars for MoE layer.") - new_tree["DeepSeekMoeBlock_0"] = {} - elif not isinstance(new_tree["DeepSeekMoeBlock_0"], dict): - max_logging.log(f"Error: 'DeepSeekMoeBlock_0' is not a dict, type: {type(new_tree['DeepSeekMoeBlock_0'])}") - new_tree["DeepSeekMoeBlock_0"] = {} - - # Merge einsum_items into the DeepSeekMoeBlock_0 dict - new_tree["DeepSeekMoeBlock_0"].update(einsum_items) - - return new_tree - - -class LayerwiseQuantization: - """ - Layerwise quantization for large models. - """ - - def __init__(self, config: Any, rng: PRNGKeyType): - self.config = config - self.rng = rng - - # The Linen path runs layer-by-layer (memory-efficient for big DeepSeek - # models) and is DeepSeek-specific because it relies on the per-layer - # `DeepSeek*ToLinen` wrappers. The NNX path runs whole-model convert - # forward and is model-agnostic — see `_load_and_quantize_nnx`. - if not config.pure_nnx: - assert config.decoder_block == common_types.DecoderBlockType.DEEPSEEK, ( - f"Linen layerwise quantization only supports {common_types.DecoderBlockType.DEEPSEEK}, " - f"got {config.decoder_block}." - ) - # Mesh definition - devices_array = maxtext_utils.create_device_mesh(config=config) - self._mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) - - self.quant = quantizations.configure_quantization(config) - if config.pure_nnx: - # NNX takes a separate code path that builds the model via from_pretrained; - # no Linen abstract-state bookkeeping is needed here. - self.unboxed_abstract_state = None - return - model = models.transformer_as_linen( - config, mesh=self._mesh, quant=self.quant, model_mode=common_types.MODEL_MODE_TRAIN - ) - init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, None, self.config, False, self.rng) - - self.unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(self.config, self._mesh, init_state_fn, False) - - def load_and_quantize(self) -> None: - """ - Load parameters layer by layer and quantize them. - """ - if self.config.pure_nnx: - self._load_and_quantize_nnx() - return - quantized_params = {} - quantized_params["params"] = {"decoder": {}} - quantized_params["aqt"] = {"decoder": {}} - config = self.config - self.quant.quant_mode = quantizations.get_quant_mode("convert") - model_mode = common_types.MODEL_MODE_PREFILL - _, rng_quant_params = jax.random.split(self.rng) - - layers = [ - deepseek.DeepSeekDenseLayerToLinen( - config=config, mesh=self._mesh, quant=self.quant, model_mode=model_mode, rngs=nnx.Rngs(self.rng) - ), - deepseek.DeepSeekMoELayerToLinen( - config=config, mesh=self._mesh, quant=self.quant, model_mode=model_mode, rngs=nnx.Rngs(self.rng) - ), - ] - layer_prefixes = [ - "dense_layers", - "moe_layers", - ] - num_moe_layers = config.num_decoder_layers - config.first_num_dense_layers - num_layers_list = [ - config.first_num_dense_layers, - num_moe_layers, - ] - - def model_apply(_p, _rng, layer): - return layer.apply( - _p | {"aqt": {}}, - jnp.ones((1, self.config.max_prefill_predict_length, self.config.base_emb_dim), dtype=jnp.int32), - None, - jnp.zeros((1, self.config.max_prefill_predict_length), dtype=jnp.int32), - True, - model_mode=model_mode, - rngs={"params": _rng}, - mutable=True, - ) - - for layer, num_layers, layer_prefix in zip(layers, num_layers_list, layer_prefixes): - for index in tqdm(range(num_layers)): - layer_name = f"{layer_prefix}_{index}" - max_logging.log(f"Processing layer: {layer_name}") - - params = self._load_layer(layer_name) - params["params"] = params["params"]["decoder"][layer_name] - - _, new_vars = model_apply(params, rng_quant_params, layer) - - if "aqt" not in new_vars: - max_logging.log( - f"Warning: 'aqt' not found in new_vars for {layer_name}. Skipping AQT processing for this layer." - ) - quantized_params["params"]["decoder"][layer_name] = params["params"] # Keep original params - continue - - aqt_vars = new_vars["aqt"] - - try: - removed_params = remove_quantized_params(params["params"], aqt_vars) - quantized_params["params"]["decoder"][layer_name] = removed_params - except Exception as e: - max_logging.log(f"ERROR: Failed to remove quantized params for {layer_name}: {e}") - max_logging.log(f"Dumping params['params'] keys for {layer_name}:") - jax.tree_util.tree_map_with_path( - lambda path, _: max_logging.log(f" {jax.tree_util.keystr(path)}"), params["params"] - ) - max_logging.log(f"Dumping new_vars['aqt'] keys for {layer_name}:") - jax.tree_util.tree_map_with_path(lambda path, _: max_logging.log(f" {jax.tree_util.keystr(path)}"), aqt_vars) - raise - - # Restructure the aqt_vars for this layer to match pure Linen format for saving - if layer_prefix == "moe_layers": - structured_aqt = insert_deepseekmoeblock_scope(aqt_vars) - else: - structured_aqt = aqt_vars - quantized_params["aqt"]["decoder"][layer_name] = structured_aqt - - unquantized_layers = ["decoder_norm", "logits_dense"] - for unquantized_layer in unquantized_layers: - params = self._load_layer(unquantized_layer) - quantized_params["params"]["decoder"][unquantized_layer] = params["params"]["decoder"][unquantized_layer] - quantized_params["params"]["token_embedder"] = self._load_layer("token_embedder")["params"]["token_embedder"] - - maxtext_utils.save_quantized_checkpoint_if_configured(self.config, quantized_params) - - def _load_and_quantize_nnx(self) -> None: - """Whole-model NNX convert: load full-precision via TRAIN-mode `from_pretrained`, - transfer kernels into a fresh CONVERT-mode model, run a forward (the - `ToNNX(AqtDotGeneral)` bridge auto-captures `qrhs.frozen`), strip kernels at - quantized paths, and save the serve-mode-shaped state. - - Two-step load: input checkpoints are typically full-precision (no AQT state - on disk), so we can't `from_pretrained(quant_mode_str="convert")` directly — - orbax would fail to find the missing `qrhs.frozen` leaves. Instead we load - in TRAIN mode (which has only kernels), then copy them into a randomly - initialized CONVERT model that already has the AQT variables provisioned. - """ - config = self.config - # MODEL_MODE_TRAIN avoids the PREFILL/AUTOREGRESSIVE cache plumbing — AQT - # layers populate `qrhs.frozen` regardless of model_mode, so train mode is - # simpler and faster. - max_logging.log("Loading full-precision NNX checkpoint in TRAIN mode...") - with self._mesh: - train_model = model_creation_utils.from_pretrained( - config, - mesh=self._mesh, - model_mode=common_types.MODEL_MODE_TRAIN, - quant_mode_str="train", - ) - - max_logging.log("Building CONVERT-mode model (random init) and copying kernels in...") - rngs = maxtext_utils_nnx.create_nnx_rngs(config, rng_key=self.rng) - with nn_partitioning.axis_rules(config.logical_axis_rules): - convert_model = model_creation_utils.from_config( - config, - mesh=self._mesh, - rngs=rngs, - model_mode=common_types.MODEL_MODE_TRAIN, - quant_mode_str="convert", - ) - self._copy_kernel_leaves_(convert_model, train_model) - del train_model - - # Forward populates AqtDotGeneral_0.qrhs.frozen on every quantized layer. - L = config.max_target_length - decoder_input_tokens = jnp.zeros((1, L), dtype=jnp.int32) - decoder_positions = jnp.arange(L, dtype=jnp.int32)[None, :] - decoder_segment_ids = jnp.ones((1, L), dtype=jnp.int32) - max_logging.log("Running CONVERT-mode forward to populate AQT scale factors...") - with nn_partitioning.axis_rules(config.logical_axis_rules): - _ = convert_model( - decoder_input_tokens, - decoder_positions, - decoder_segment_ids=decoder_segment_ids, - enable_dropout=False, - model_mode=common_types.MODEL_MODE_TRAIN, - ) - - # Convert-mode state has both `kernel` (full precision) and `AqtDotGeneral_0.qrhs.frozen` - # at every quantized DenseGeneral; the serve-mode reader expects only the latter. - convert_state = nnx.state(convert_model).to_pure_dict() - serve_state = self._strip_kernels_at_quantized_paths(convert_state) - - if config.save_quantized_params_path: - max_logging.log(f"Saving NNX-format quantized checkpoint to {config.save_quantized_params_path}") - - # Wrap each leaf in `{"value": }` so the on-disk shape matches what - # `from_pretrained`'s NNX-detection branch reads back (it later does - # `tree.map(lambda v: v["value"], ...)` on each leaf). Save directly via - # orbax — `save_params_to_path` would add an outer `{"params": ...}` wrap - # that the NNX path doesn't expect. - def _wrap_value(node): - if isinstance(node, dict): - return {k: _wrap_value(v) for k, v in node.items()} - return {"value": node} - - wrapped = _wrap_value(serve_state) - orbax_checkpointer = ocp.PyTreeCheckpointer( - use_ocdbt=config.checkpoint_storage_use_ocdbt, - use_zarr3=config.checkpoint_storage_use_zarr3, - ) - orbax_checkpointer.save(config.save_quantized_params_path, wrapped, force=True) - max_logging.log(f"Saved NNX-format quantized checkpoint at: {config.save_quantized_params_path}") - else: - max_logging.log("Skipping save: save_quantized_params_path is null.") - - @staticmethod - def _copy_kernel_leaves_(dst_model, src_model): - """Copy the full-precision parameter leaves (kernel/embedding/scale/bias) - from src into dst, leaving dst's AQT and RNG variables untouched. - """ - src_dict = nnx.state(src_model).to_pure_dict() - dst_state = nnx.state(dst_model) - dst_dict = dst_state.to_pure_dict() - - def walk(d_node, s_node): - if not (isinstance(d_node, dict) and isinstance(s_node, dict)): - return - for key, d_child in d_node.items(): - if key not in s_node: - continue - s_child = s_node[key] - if key in ("kernel", "embedding", "scale", "bias") and not isinstance(d_child, dict): - d_node[key] = s_child - elif isinstance(d_child, dict): - walk(d_child, s_child) - - walk(dst_dict, src_dict) - nnx.replace_by_pure_dict(dst_state, dst_dict) - nnx.update(dst_model, dst_state) - - @staticmethod - def _strip_kernels_at_quantized_paths(state_dict): - """Drop `kernel` keys at any node that has a sibling `AqtDotGeneral_0`. - - In convert mode each quantized DenseGeneral keeps both the full-precision - `kernel` (an nnx.Param) and the AQT-quantized `AqtDotGeneral_0.qrhs.frozen` - side-by-side. Serve mode (the on-disk shape `from_pretrained` reads back) - only carries the latter; the kernel is recreated as a dummy zero in - `linears.DenseGeneral.__call__`. - """ - if not isinstance(state_dict, dict): - return state_dict - has_aqt = "AqtDotGeneral_0" in state_dict - out = {} - for k, v in state_dict.items(): - if k == "kernel" and has_aqt: - continue - out[k] = LayerwiseQuantization._strip_kernels_at_quantized_paths(v) if isinstance(v, dict) else v - return out - - def _load_layer(self, layer_name): - """Loads a specific layer's parameters from the checkpoint.""" - - config = self.config - with nn_partitioning.axis_rules(config.logical_axis_rules): - - params = checkpointing.load_params_from_path( - config.load_parameters_path, - self._create_partial_abstract_params(self.unboxed_abstract_state.params, layer_name), - config.checkpoint_storage_concurrent_gb, - config.checkpoint_storage_use_ocdbt, - config.checkpoint_storage_use_zarr3, - ) - return params - - def _create_partial_abstract_params(self, abstract_unboxed_params, layer): - """Creates a partial abstract params structure using ocp.PLACEHOLDER.""" - - def _should_keep(path, _): - # True if the layer name is part of the path - return any(isinstance(key, jax.tree_util.DictKey) and key.key == layer for key in path) - - def _map_fn(path, value): - if not _should_keep(path, value): - return IGNORE - if isinstance(value, jax.ShapeDtypeStruct): - zeros_array = jnp.zeros(value.shape, value.dtype) - if value.sharding is not None: - try: - return jax.device_put(zeros_array, value.sharding) - except Exception as e: # pylint: disable=broad-except - max_logging.log(f"Error applying sharding for path {path}: {e}") - return zeros_array - return zeros_array - return value - - return jax.tree_util.tree_map_with_path(_map_fn, abstract_unboxed_params) - - -def main(argv: Sequence[str]) -> None: - jax.config.update("jax_default_prng_impl", "unsafe_rbg") - os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0" - config = pyconfig.initialize(argv) - validate_config(config) - max_utils.print_system_information() - rng = jax.random.PRNGKey(1234) - quantization = LayerwiseQuantization(config, rng) - quantization.load_and_quantize() - - -def validate_config(config): - assert ( - config.load_full_state_path == "" - ), "Operation on full states not supported! Convert to parameter checkpoint first." - - -if __name__ == "__main__": - app.run(main) diff --git a/tests/__init__.py b/tests/__init__.py index 46cd7ffa11..76d26765cb 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -16,6 +16,54 @@ Test initialization """ -import pathwaysutils +try: + import pathwaysutils -pathwaysutils.initialize() + pathwaysutils.initialize() +except ImportError: + import sys + from unittest.mock import MagicMock + + mock_pathwaysutils = MagicMock() + mock_pathwaysutils.__path__ = [] + mock_pathwaysutils.is_pathways_backend_used.return_value = False + sys.modules["pathwaysutils"] = mock_pathwaysutils + + mock_elastic = MagicMock() + mock_elastic.__path__ = [] + sys.modules["pathwaysutils.elastic"] = mock_elastic + + mock_manager = MagicMock() + sys.modules["pathwaysutils.elastic.manager"] = mock_manager + +try: + import tokamax + from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_kernel +except ImportError: + import sys + from unittest.mock import MagicMock + + mock_tokamax = MagicMock() + mock_tokamax.__path__ = [] + sys.modules["tokamax"] = mock_tokamax + sys.modules["tokamax._src"] = MagicMock() + sys.modules["tokamax._src.ops"] = MagicMock() + sys.modules["tokamax._src.ops.experimental"] = MagicMock() + sys.modules["tokamax._src.ops.experimental.tpu"] = MagicMock() + sys.modules["tokamax._src.ops.experimental.tpu.splash_attention"] = MagicMock() + +try: + import tensorflow +except ImportError: + import sys + from unittest.mock import MagicMock + + mock_tf = MagicMock() + mock_tf.__path__ = [] + sys.modules["tensorflow"] = mock_tf + sys.modules["tensorflow.io"] = MagicMock() + sys.modules["tensorflow.data"] = MagicMock() + sys.modules["tensorflow.compat"] = MagicMock() + sys.modules["tensorflow.compat.v1"] = MagicMock() + sys.modules["tensorflow.compat.v1.io"] = MagicMock() + sys.modules["tensorflow.compat.v1.io.gfile"] = MagicMock() diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index b0af64d9fc..192cf82867 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -14,313 +14,74 @@ """Tests for the quantizations""" -import functools -import os.path import sys -from typing import Any import unittest -from aqt.jax.v2 import aqt_tensor -from aqt.jax.v2.flax import aqt_flax -from flax import nnx import jax -from jax import lax from jax import numpy as jnp from jax.sharding import Mesh from maxtext.configs import pyconfig -from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR -from maxtext.kernels.megablox import gmm -from maxtext.layers import nnx_wrappers, quantizations +from flax import nnx +from maxtext.layers import moe +from maxtext.layers import quantizations +from maxtext.layers.initializers import nd_dense_init from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils from tests.utils.test_helpers import get_test_config_path -import numpy as np import pytest -_QUERY_REGEX = ".*/query" -_VALUE_REGEX = ".*/value" - - -class QuantTestModule(nnx.Module): - """Test module for einsum.""" - - def __init__( - self, - quantization: quantizations.AqtQuantization, - data_type: Any, - rngs: nnx.Rngs, - ): - self.quantization = quantization - self.identity = jnp.identity(2, dtype=data_type) - self.einsum = None - self.dot_general = None - - if self.quantization: - quant_dg, is_tiled, tiling_fn = None, False, None - if isinstance(self.quantization.quant_dg, dict): - quant_dg, is_tiled, tiling_fn = self._get_mixed_precision_cfg() - else: - quant_dg, is_tiled, tiling_fn = self.quantization.quant_dg, False, None - rhs_axis_metadata_wrapper = None - if self.quantization.quant_mode == aqt_flax.QuantMode.CONVERT: - rhs_axis_metadata_wrapper = None - else: - rhs_axis_metadata_wrapper = functools.partial( - quantizations._rhs_axis_metadata_wrapper, - mesh_axes=(), - is_tiled=is_tiled, - replicate_scale=self.quantization.replicate_scale, - ) - - aqt_dg_cls = aqt_flax.AqtDotGeneral( - quant_dg, - rhs_quant_mode=self.quantization.quant_mode, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.CALIBRATION_AND_VALUE, - rhs_axis_metadata_wrapper=rhs_axis_metadata_wrapper, - use_legacy_freezer=False, - tiling_fn=tiling_fn, - ) - aqt_dg_cls_nnx = nnx_wrappers.ToNNX(aqt_dg_cls, rngs=nnx.Rngs(params=0)) - aqt_einsum = aqt_flax.AqtEinsum( - cfg=quant_dg, - rhs_quant_mode=self.quantization.quant_mode, - lhs_freeze_mode=aqt_flax.FreezerMode.NONE, - rhs_freeze_mode=aqt_flax.FreezerMode.CALIBRATION_AND_VALUE, - rhs_axis_metadata_wrapper=rhs_axis_metadata_wrapper, - use_legacy_freezer=False, - tiling_fn=tiling_fn, - ) - aqt_einsum_nnx = nnx_wrappers.ToNNX(aqt_einsum, rngs=nnx.Rngs(params=0)) - self.einsum = nnx.data(aqt_einsum_nnx) - self.dot_general = nnx.data(aqt_dg_cls_nnx) - else: - self.einsum = jnp.einsum - self.dot_general = lax.dot_general - - def __call__(self, inputs): - res_einsum = self.einsum("bc,ab->ac", inputs, self.identity) - res_dg = self.dot_general(inputs, inputs, (((), ()), ((), ())), precision=None) - return res_einsum, res_dg - - -def _configure_quantization(quant_str="", quant_cfg_path="", mode_str="train", replicate_scale=False): +def _configure_quantization(quant_str="", mode_str="train"): config = pyconfig.initialize( [None, get_test_config_path()], enable_checkpointing=False, quantization=quant_str, - quant_cfg_path=quant_cfg_path, - replicate_quant_scale=replicate_scale, ) quant = quantizations.configure_quantization(config, mode_str) return quant -def _apply(quant_str=""): - rngs = nnx.Rngs(params=0) - inputs = jnp.ones((2, 2)) - data_type = inputs.dtype - quant = _configure_quantization(quant_str) - test_module = QuantTestModule(quant, data_type, rngs) - res_einsum, res_dg = test_module(inputs) - return inputs, res_einsum, res_dg - - class QuantizationTest(unittest.TestCase): """Tests for quantization.""" - def test_in_quant_mode(self): - quant = _configure_quantization(quant_str="int8", mode_str="convert") - self.assertTrue(quantizations.in_convert_mode(quant)) - self.assertFalse(quantizations.in_serve_mode(quant)) - def test_configure_quantization_is_null(self): for quant_mode in ["train", "serve", "convert"]: quant = _configure_quantization(quant_str="", mode_str=quant_mode) self.assertEqual(quant, None) - def test_configure_quantization_replicate_scale(self): - for quant_mode in ["train", "serve", "convert"]: - quant = _configure_quantization(quant_str="int8", mode_str=quant_mode) - self.assertEqual(quant.replicate_scale, False) - for quant_mode in ["train", "serve", "convert"]: - quant = _configure_quantization(quant_str="int8", mode_str=quant_mode, replicate_scale=True) - self.assertEqual(quant.replicate_scale, True) +class QuantizationConfigValidationTest(unittest.TestCase): + """Tests for quantization configuration validation.""" - @pytest.mark.cpu_only - def test_configure_quantization_is_int8(self): - for quant_mode in ["train", "serve", "convert"]: - quant = _configure_quantization(quant_str="int8", mode_str=quant_mode) - self.assertNotEqual(quant, None) - - @pytest.mark.tpu_only # b/421002974 - def test_aqt_quantization(self): - # Without quantization - inputs, res_einsum, res_dg = _apply() - self.assertTrue(jnp.array_equal(inputs, res_einsum)) - self.assertEqual(res_einsum.dtype, np.dtype(np.float32)) - self.assertTrue(jnp.array_equal(inputs, res_dg[0][0])) - self.assertEqual(res_dg.dtype, np.dtype(np.float32)) - - # With int8 quantization - inputs, res_einsum, res_dg = _apply(quant_str="int8") - self.assertTrue(jnp.greater(jnp.max(inputs), jnp.max(res_einsum))) - self.assertEqual(res_einsum.dtype, np.dtype(np.float32)) - self.assertTrue(jnp.greater(jnp.max(inputs), jnp.max(res_dg[0][0]))) - # self.assertEqual(res_dg.dtype, np.dtype(np.float32)) - - def test_mixed_precision_config_int8w(self): - quant = _configure_quantization( - quant_str="intmp", - quant_cfg_path=os.path.join(MAXTEXT_CONFIGS_DIR, "quantization", "int8_weight_only.json"), - ) - self.assertTrue(isinstance(quant.quant_dg, dict) and len(quant.quant_dg) == 1) - # pylint: disable=unsupported-membership-test - self.assertTrue(quantizations.DEFAULT in quant.quant_dg) - quant_cfg, _ = quant.quant_dg[quantizations.DEFAULT] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.dtype, None) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 8) - - def test_mixed_precision_config_scale(self): - quant = _configure_quantization( - quant_str="intmp", - quant_cfg_path=os.path.join( - MAXTEXT_CONFIGS_DIR, - "quantization", - "dense_llm_weight_only_scale.json", - ), - ) - self.assertTrue(isinstance(quant.quant_dg, dict) and len(quant.quant_dg) == 7) - # pylint: disable=unsupported-membership-test - self.assertTrue(quantizations.DEFAULT in quant.quant_dg) - quant_cfg, _ = quant.quant_dg[quantizations.DEFAULT] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.dtype, None) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 8) - quant_cfg, _ = quant.quant_dg[_QUERY_REGEX] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.dtype, None) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 4) - - def test_mixed_precision_config_subchannel(self): - quant = _configure_quantization( - quant_str="intmp", - quant_cfg_path=os.path.join( - MAXTEXT_CONFIGS_DIR, - "quantization", - "dense_llm_subchannel.json", - ), + def test_use_qwix_quantization_default(self): + # Verify that use_qwix_quantization defaults to True when initializing config + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="int8", ) - self.assertTrue(isinstance(quant.quant_dg, dict) and len(quant.quant_dg) == 7) - # pylint: disable=unsupported-membership-test - self.assertTrue(quantizations.DEFAULT in quant.quant_dg) - quant_cfg, tile_size = quant.quant_dg[quantizations.DEFAULT] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.bits, 8) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 8) - self.assertEqual(tile_size, -1) - quant_cfg, tile_size = quant.quant_dg[_QUERY_REGEX] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.bits, 8) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 4) - self.assertEqual(tile_size, 128) - - quant_cfg, tile_size = quant.quant_dg[_VALUE_REGEX] - self.assertEqual(quant_cfg.fwd.dg_quantizer.lhs.numerics.bits, 8) - self.assertEqual(quant_cfg.fwd.dg_quantizer.rhs.numerics.bits, 4) - self.assertEqual(tile_size, -1) - - def test_remove_quantized_params(self): - _params = { - "decoder": { - "decoder_norm": {"scale": 1.0}, - "layers": { - "mlp": { - "wi_0": {"kernel": 1.0}, - "wi_1": {"kernel": 1.0}, - "wo": {"kernel": 1.0}, - }, - "self_attention": { - "key": {"kernel": 1.0}, - }, - }, - "logits_dense": {"kernel": 1.0}, - }, - } - _aqt_vars = { - "decoder": { - "layers": { - "mlp": { - "wi_0": { - "AqtDotGeneral_0": { - "qrhs": { - "frozen": aqt_tensor.QTensor( - qvalue=[1.1, 1.0], - scale=[1.0], - scale_t=[1.0], - bias=1.0, - ) - } - } - }, - "wi_1": { - "AqtDotGeneral_0": { - "qrhs": { - "frozen": aqt_tensor.QTensor( - qvalue=[1.1, 1.0], - scale=[1.0], - scale_t=[1.0], - bias=1.0, - ) - } - } - }, - "wo": { - "AqtDotGeneral_0": { - "qrhs": { - "frozen": aqt_tensor.QTensor( - qvalue=[1.1, 1.0], - scale=[1.0], - scale_t=[1.0], - bias=1.0, - ) - } - } - }, - }, - "self_attention": { - "key": { - "AqtDotGeneral_0": { - "qrhs": { - "frozen": aqt_tensor.QTensor( - qvalue=[1.1, 1.0], - scale=[1.0], - scale_t=[1.0], - bias=1.0, - ) - } - } - } - }, - } - } - } - _expected = { - "decoder": { - "decoder_norm": {"scale": 1.0}, - "layers": { - "mlp": { - "wi_0": {"kernel": {}}, - "wi_1": {"kernel": {}}, - "wo": {"kernel": {}}, - }, - "self_attention": { - "key": {"kernel": {}}, - }, - }, - "logits_dense": {"kernel": 1.0}, - } - } - result = quantizations.remove_quantized_params(_params, _aqt_vars) - self.assertEqual(_expected, result) + self.assertTrue(config.use_qwix_quantization) + + def test_unsupported_quantization_without_qwix(self): + # Verify that setting use_qwix_quantization=False with non-native/non-TE quantization raises ValueError + with self.assertRaisesRegex(ValueError, "is unsupported because legacy AQT has been completely removed"): + pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="int8", + use_qwix_quantization=False, + ) + + def test_supported_quantization_without_qwix(self): + # Verify that setting use_qwix_quantization=False with native FP8 or TE is allowed and does not raise + for quant_type in ["fp8", "nanoo_fp8", "fp8_gpu", "te_fp8_delayedscaling"]: + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization=quant_type, + use_qwix_quantization=False, + ) + self.assertFalse(config.use_qwix_quantization) class QuantTest(unittest.TestCase): @@ -524,6 +285,7 @@ def test_fp8_te_nvfp4_quantization(self): ) @pytest.mark.tpu_only def test_gmm_kernel(group_sizes, k, n, tiling, dtype): + # pylint: disable=undefined-variable """Smoke-test + correctness check for the grouped matrix-multiply kernel. For each group i, gmm should compute @@ -564,5 +326,119 @@ def test_gmm_kernel(group_sizes, k, n, tiling, dtype): assert jnp.abs(quant_out - base_out).mean() / jnp.abs(base_out).mean() < 2e-1 +class QuantizationCoverageTest(unittest.TestCase): + """Explicit tests to ensure 100% test coverage of all quantization paths.""" + + def test_configure_quantization_paths(self): + # Test all configure_quantization paths on CPU (instantiation only) + config_fp8 = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="fp8", + use_qwix_quantization=False, + ) + quant_fp8 = quantizations.configure_quantization(config_fp8, "train") + self.assertIsNotNone(quant_fp8) + + config_nanoo = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="nanoo_fp8", + use_qwix_quantization=False, + ) + quant_nanoo = quantizations.configure_quantization(config_nanoo, "train") + self.assertIsNotNone(quant_nanoo) + + # Only run TE quantization config test if transformer_engine is installed + try: + import transformer_engine # pylint: disable=unused-import,import-outside-toplevel + + has_te = True + except ImportError: + has_te = False + + if has_te: + config_te = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="te_fp8_delayedscaling", + use_qwix_quantization=False, + ) + quant_te = quantizations.configure_quantization(config_te, "train") + self.assertIsNotNone(quant_te) + + def test_configure_kv_quant(self): + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantize_kvcache=False, + ) + # Should not raise + quantizations.configure_kv_quant(config) + + config_fail = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantize_kvcache=True, + ) + with self.assertRaises(ValueError): + quantizations.configure_kv_quant(config_fail) + + def test_moe_quantization_coverage(self): + # Instantiates RoutedMoE on CPU to cover the AQT-free parameter initialization path in moe.py + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + quantization="int8", + use_qwix_quantization=True, + num_experts=2, + base_emb_dim=8, + base_mlp_dim=8, + base_moe_mlp_dim=8, # Required positive base value to derive positive moe_mlp_dim + parameter_memory_host_offload=True, # Cover the parameter offloading paths in linears.py + ) + + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + + quant = quantizations.configure_quantization(config, "train") + + with mesh: + moe_layer = moe.RoutedMoE( + config=config, + num_experts=config.num_experts, + num_experts_per_tok=1, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("expert", "embed_moe", "heads"), + rngs=nnx.Rngs(0), + quant=quant, + ) + + # In Flax NNX, parameters are fully initialized during instantiation. + self.assertIsNotNone(moe_layer.gate.kernel) + + # Execute a forward pass to cover DenseGeneral.__call__, RoutedMoE.__call__, + # sparse_matmul, and the custom quant_einsum wrapper in moe.py + inputs = jnp.ones((2, 4, 8), dtype=jnp.float32) + outputs, _, _ = moe_layer(inputs) + self.assertEqual(outputs.shape, (2, 4, 8)) + + def test_quantization_fallbacks(self): + # Cover the fallback return None path in _get_quant_config when an unsupported scheme is passed + config_invalid = pyconfig.initialize( + [None, get_test_config_path()], + quantization="int4", + ) + self.assertIsNone(quantizations.configure_quantization(config_invalid)) + + # Cover the implicit return None path in configure_kv_quant when quantize_kvcache is False + config_no_kv = pyconfig.initialize( + [None, get_test_config_path()], + quantize_kvcache=False, + ) + self.assertIsNone(quantizations.configure_kv_quant(config_no_kv)) + + if __name__ == "__main__": unittest.main() From 79c978bccfc138beeca5284067b9f37470d9c699 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Fri, 5 Jun 2026 22:30:03 +0000 Subject: [PATCH 02/52] Fix ShardingTypeError in test_explicit_shard_mode by reverting out_sharding to sharding in RoutedMoE --- src/maxtext/layers/moe.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index e222831388..4450d60c31 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -485,7 +485,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wi_kernel_axes, + sharding=self.wi_kernel_axes, ) self.wo = nnx.Param( self.kernel_init( @@ -495,7 +495,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wo_kernel_axes, + sharding=self.wo_kernel_axes, ) else: # Pad model dimension in Unfused MoE weight kernels for GMM_v2 execution. @@ -512,7 +512,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wi_kernel_axes, + sharding=self.wi_kernel_axes, ) self.wi_1 = nnx.Param( self.kernel_init( @@ -522,7 +522,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wi_kernel_axes, + sharding=self.wi_kernel_axes, ) self.wo = nnx.Param( self.kernel_init( @@ -532,7 +532,7 @@ def __init__( kernel_in_axis, kernel_out_axis, ), - out_sharding=self.wo_kernel_axes, + sharding=self.wo_kernel_axes, ) if self.config.mlp_bias: @@ -542,15 +542,15 @@ def __init__( wo_bias_shape = (self.num_experts, self.moe_expert_input_dim) self.wi_0_bias = nnx.Param( default_bias_init(self.rngs.params(), wi_bias_shape, self.weight_dtype), - out_sharding=wi_bias_axes, + sharding=wi_bias_axes, ) self.wi_1_bias = nnx.Param( default_bias_init(self.rngs.params(), wi_bias_shape, self.weight_dtype), - out_sharding=wi_bias_axes, + sharding=wi_bias_axes, ) self.wo_bias = nnx.Param( default_bias_init(self.rngs.params(), wo_bias_shape, self.weight_dtype), - out_sharding=wo_bias_axes, + sharding=wo_bias_axes, ) else: self.wi_0_bias = None @@ -560,7 +560,7 @@ def __init__( if self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4: self.per_expert_scale = nnx.Param( jnp.ones((self.num_experts,), dtype=self.weight_dtype), - out_sharding=("exp",), + sharding=("exp",), ) else: self.per_expert_scale = None From 290725f69bb519239beadc7547bb4ca88fa5c6f0 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 00:21:20 +0000 Subject: [PATCH 03/52] Fix MoE, quantization, distillation, and MaxEngine test failures --- src/maxtext/inference/maxengine/maxengine.py | 12 +++--------- src/maxtext/layers/moe.py | 2 +- src/maxtext/layers/quantizations.py | 2 +- tests/integration/decode_tests.py | 1 - tests/unit/quantizations_test.py | 1 + 5 files changed, 6 insertions(+), 12 deletions(-) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 74709cd28a..ed0b5ff74a 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -383,15 +383,9 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar return params def _load_params_nnx(self, params, rng): - """NNX equivalent of load_params: returns an nnx.Param state and populates KV cache shardings. - - Quantization handling: - * `checkpoint_is_quantized=True`: model built in `serve` mode (no full - kernel), `from_pretrained` reads `qrhs.frozen` from disk. - * `checkpoint_is_quantized=False` + `quantization=...`: model built in - `train` mode, full-precision kernel loaded; AQT layers quantize per - forward. Same output as serve mode (absmax calibration), slower. - """ + """NNX equivalent of load_params: returns an nnx.Param state and populates KV cache shardings.""" + if self.config.quantization: + raise NotImplementedError("pure_nnx + quantization not yet supported. Use pure_nnx=False.") if params: print("Resharding given NNX params") diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 4450d60c31..c29cb4eeee 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -2339,7 +2339,7 @@ def __call__( output, lb_loss, bias_updates = self.fused_moe_matmul( inputs, gate_logits, wo_kernel, w0_kernel=w0_kernel, w1_kernel=w1_kernel, fused_kernel=fused_kernel ) - + elif cfg.sparse_matmul: output, lb_loss, bias_updates = self.sparse_matmul( inputs, gate_logits, pre_bias_logits, w0_kernel, w1_kernel, wo_kernel, w0_bias, w1_bias, wo_bias ) diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 6917e27127..f7f64f7385 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -386,7 +386,7 @@ def get_qt_provider(config): def maybe_quantize_model(model, config): """Quantize the model if quantization is enabled.""" # Batch split is not using Qwix's interception feature but manual plumbing - if config.use_qwix_quantization and not config.use_batch_split_schedule: + if config.use_qwix_quantization and not config.use_batch_split_schedule and not config.pure_nnx: quantization_provider = get_qt_provider(config) if quantization_provider: model = qwix.quantize_model(model, quantization_provider) diff --git a/tests/integration/decode_tests.py b/tests/integration/decode_tests.py index 0117dc1a6b..3af2f6d03a 100644 --- a/tests/integration/decode_tests.py +++ b/tests/integration/decode_tests.py @@ -62,7 +62,6 @@ class DecodeTests(unittest.TestCase): "max_target_length=128", "per_device_batch_size=1", "quantization=int8", - "quantize_kvcache=True", rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", ], "pdb_lt_1": [ # tests decode with per_device_batch_size < 1 diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 192cf82867..aca57801aa 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -24,6 +24,7 @@ from flax import nnx from maxtext.layers import moe from maxtext.layers import quantizations +from maxtext.kernels.megablox.ops import gmm from maxtext.layers.initializers import nd_dense_init from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils From e98f7ade288cc870c8036f7c58018403562b6005 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 00:33:34 +0000 Subject: [PATCH 04/52] Add unit tests to cover missing/partial lines in linears, quantizations, and moe --- tests/unit/quantizations_test.py | 69 ++++++++++++++++++++++++++++++++ 1 file changed, 69 insertions(+) diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index aca57801aa..3238337a91 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -440,6 +440,75 @@ def test_quantization_fallbacks(self): ) self.assertIsNone(quantizations.configure_kv_quant(config_no_kv)) + def test_dense_general_parameter_offload_coverage(self): + # Covers parameter_memory_host_offload paths in linears.py + from maxtext.layers import linears + + dense_layer = linears.DenseGeneral( + in_features_shape=8, + out_features_shape=8, + parameter_memory_host_offload=True, + rngs=nnx.Rngs(0), + ) + inputs = jnp.ones((2, 8), dtype=jnp.float32) + outputs = dense_layer(inputs) + self.assertEqual(outputs.shape, (2, 8)) + + def test_configure_quantization_batch_split_schedule(self): + # Covers use_batch_split_schedule path in quantizations.py + config_bs = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + use_batch_split_schedule=True, + quantization="fp8_full", + use_qwix_quantization=False, + ) + quant = quantizations.configure_quantization(config_bs, "train") + self.assertIsInstance(quant, quantizations.QwixQuantization) + + config_bs_manual = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + use_batch_split_schedule=True, + quantization="fp8_full", + use_manual_quantization=True, + use_qwix_quantization=False, + ) + quant_manual = quantizations.configure_quantization(config_bs_manual, "train") + self.assertIsNone(quant_manual) + + def test_moe_gemma4_coverage(self): + # Covers GEMMA4 routing and expert scale fusion paths in moe.py + from maxtext.layers import moe + + config = pyconfig.initialize( + [None, get_test_config_path()], + enable_checkpointing=False, + decoder_block="gemma4", + model_call_mode="inference", + fuse_expert_scales=True, + num_experts=2, + base_emb_dim=8, + base_mlp_dim=8, + base_moe_mlp_dim=8, + ) + devices_array = maxtext_utils.create_device_mesh(config) + mesh = Mesh(devices_array, config.mesh_axes) + + with mesh: + moe_layer = moe.RoutedMoE( + config=config, + num_experts=config.num_experts, + num_experts_per_tok=1, + mesh=mesh, + kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"), + kernel_axes=("expert", "embed_moe", "heads"), + rngs=nnx.Rngs(0), + ) + inputs = jnp.ones((2, 4, 8), dtype=jnp.float32) + outputs, _, _ = moe_layer(inputs) + self.assertEqual(outputs.shape, (2, 4, 8)) + if __name__ == "__main__": unittest.main() From 0c9be00b268c765d16888b87b2ccd1393206af09 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 01:03:19 +0000 Subject: [PATCH 05/52] Change use_qwix_quantization to True in test_configure_quantization_paths --- tests/unit/quantizations_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 3238337a91..662cd23319 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -336,7 +336,7 @@ def test_configure_quantization_paths(self): [None, get_test_config_path()], enable_checkpointing=False, quantization="fp8", - use_qwix_quantization=False, + use_qwix_quantization=True, ) quant_fp8 = quantizations.configure_quantization(config_fp8, "train") self.assertIsNotNone(quant_fp8) @@ -345,7 +345,7 @@ def test_configure_quantization_paths(self): [None, get_test_config_path()], enable_checkpointing=False, quantization="nanoo_fp8", - use_qwix_quantization=False, + use_qwix_quantization=True, ) quant_nanoo = quantizations.configure_quantization(config_nanoo, "train") self.assertIsNotNone(quant_nanoo) @@ -363,7 +363,7 @@ def test_configure_quantization_paths(self): [None, get_test_config_path()], enable_checkpointing=False, quantization="te_fp8_delayedscaling", - use_qwix_quantization=False, + use_qwix_quantization=True, ) quant_te = quantizations.configure_quantization(config_te, "train") self.assertIsNotNone(quant_te) From e23ce496a58a56981c7ce79306dfaf0c06585bdc Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 01:48:19 +0000 Subject: [PATCH 06/52] Update HLO references for Qwix/FP8 transition --- tests/utils/reference_hlo_deepseek3.txt | 1948 +++++++++---------- tests/utils/reference_hlo_llama3_8b.txt | 2168 ++++++++++----------- tests/utils/reference_hlo_qwen3_1.7b.txt | 2212 +++++++++++----------- 3 files changed, 3164 insertions(+), 3164 deletions(-) diff --git a/tests/utils/reference_hlo_deepseek3.txt b/tests/utils/reference_hlo_deepseek3.txt index 39c44fd147..de9a31f4e5 100644 --- a/tests/utils/reference_hlo_deepseek3.txt +++ b/tests/utils/reference_hlo_deepseek3.txt @@ -10,21 +10,21 @@ StackFrames %region_46.56 (top_k.25: bf16[], top_k.26: bf16[], top_k.27: s32[], top_k.28: s32[]) -> pred[] { - %constant.1424 = s32[]{:T(128)} constant(0) - %constant.1425 = s32[]{:T(128)} constant(2147483647) + %constant.1408 = s32[]{:T(128)} constant(0) + %constant.1409 = s32[]{:T(128)} constant(2147483647) %top_k.25 = bf16[]{:T(256)} parameter(0), metadata={op_name="top_k"} %top_k.26 = bf16[]{:T(256)} parameter(1), metadata={op_name="top_k"} %top_k.27 = s32[]{:T(128)} parameter(2), metadata={op_name="top_k"} %top_k.28 = s32[]{:T(128)} parameter(3), metadata={op_name="top_k"} %convert.393 = f32[]{:T(128)S(6)} convert(%top_k.25), metadata={op_name="convert.18"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %bitcast-convert.39 = s32[]{:T(128)S(6)} bitcast-convert(%convert.393), metadata={op_name="bitcast-convert.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.144 = pred[]{:T(512)S(6)} compare(%bitcast-convert.39, %constant.1424), direction=LT, metadata={op_name="compare.38"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.40 = s32[]{:T(128)S(6)} xor(%constant.1425, %bitcast-convert.39), metadata={op_name="xor.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.144 = pred[]{:T(512)S(6)} compare(%bitcast-convert.39, %constant.1408), direction=LT, metadata={op_name="compare.38"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.40 = s32[]{:T(128)S(6)} xor(%constant.1409, %bitcast-convert.39), metadata={op_name="xor.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.127 = s32[]{:T(128)S(6)} select(%compare.144, %xor.40, %bitcast-convert.39), metadata={op_name="select.16"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %convert.394 = f32[]{:T(128)S(6)} convert(%top_k.26), metadata={op_name="convert.19"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %bitcast-convert.40 = s32[]{:T(128)S(6)} bitcast-convert(%convert.394), metadata={op_name="bitcast-convert.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.145 = pred[]{:T(512)S(6)} compare(%bitcast-convert.40, %constant.1424), direction=LT, metadata={op_name="compare.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.41 = s32[]{:T(128)S(6)} xor(%constant.1425, %bitcast-convert.40), metadata={op_name="xor.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.145 = pred[]{:T(512)S(6)} compare(%bitcast-convert.40, %constant.1408), direction=LT, metadata={op_name="compare.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.41 = s32[]{:T(128)S(6)} xor(%constant.1409, %bitcast-convert.40), metadata={op_name="xor.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.128 = s32[]{:T(128)S(6)} select(%compare.145, %xor.41, %bitcast-convert.40), metadata={op_name="select.17"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %compare.146 = pred[]{:T(512)S(6)} compare(%select.127, %select.128), direction=GT, metadata={op_name="compare.0"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %compare.147 = pred[]{:T(512)S(6)} compare(%select.128, %select.127), direction=GT, metadata={op_name="compare.117"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} @@ -78,19 +78,19 @@ StackFrames %region_107.126 (psum.6: bf16[], psum.9: bf16[]) -> bf16[] { %psum.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} %psum.9 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1417 = bf16[]{:T(256)} add(%psum.6, %psum.9), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.1407 = bf16[]{:T(256)} add(%psum.6, %psum.9), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %region_108.127 (psum.10: bf16[], psum.11: bf16[]) -> bf16[] { %psum.10 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} %psum.11 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1418 = bf16[]{:T(256)} add(%psum.10, %psum.11), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.1408 = bf16[]{:T(256)} add(%psum.10, %psum.11), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %region_109.128 (psum.14: bf16[], psum.15: bf16[]) -> bf16[] { %psum.14 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} %psum.15 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1419 = bf16[]{:T(256)} add(%psum.14, %psum.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.1409 = bf16[]{:T(256)} add(%psum.14, %psum.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %region_62.73 (reduce-window.111: s32[], reduce-window.112: s32[]) -> s32[] { @@ -211,167 +211,167 @@ StackFrames %param_0.17 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) %param_1.108 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.13 = s32[1024]{0:T(1024)} custom-call(%param_1.108), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %slice.892 = s32[512]{0:T(512)} slice(%custom-call.13), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %reshape.3298 = s32[4,128]{1,0:T(4,128)} reshape(%slice.892), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.847 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3298), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %gather.183 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} gather(%param_0.17, %transpose.847), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %transpose.846 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%gather.183), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - ROOT %reshape.3297 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.846), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %slice.920 = s32[512]{0:T(512)} slice(%custom-call.13), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %reshape.3318 = s32[4,128]{1,0:T(4,128)} reshape(%slice.920), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %transpose.847 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3318), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %gather.187 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} gather(%param_0.17, %transpose.847), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %transpose.846 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%gather.187), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + ROOT %reshape.3317 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.846), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } %fused_computation.6 (param_0.20: f32[163840,32], param_1.110: s32[1024]) -> f32[512,32] { - %param_0.20 = f32[163840,32]{1,0:T(8,128)} parameter(0) + %param_0.20 = f32[163840,32]{1,0:T(8,128)S(1)} parameter(0) %param_1.110 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.15 = s32[1024]{0:T(1024)} custom-call(%param_1.110), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %slice.894 = s32[512]{0:T(512)} slice(%custom-call.15), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %reshape.3306 = s32[4,128]{1,0:T(4,128)} reshape(%slice.894), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %transpose.853 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3306), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %gather.185 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.20, %transpose.853), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %transpose.852 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.185), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - ROOT %reshape.3305 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.852), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %slice.922 = s32[512]{0:T(512)} slice(%custom-call.15), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %reshape.3326 = s32[4,128]{1,0:T(4,128)} reshape(%slice.922), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} + %transpose.853 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3326), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} + %gather.189 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.20, %transpose.853), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %transpose.852 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.189), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + ROOT %reshape.3325 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.852), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} } %fused_computation.7 (param_0.23: f32[163840,32], param_1.112: s32[1024]) -> f32[512,32] { %param_0.23 = f32[163840,32]{1,0:T(8,128)} parameter(0) %param_1.112 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.17 = s32[1024]{0:T(1024)} custom-call(%param_1.112), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %slice.896 = s32[512]{0:T(512)} slice(%custom-call.17), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %reshape.3314 = s32[4,128]{1,0:T(4,128)} reshape(%slice.896), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %transpose.859 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3314), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %gather.187 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.23, %transpose.859), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %transpose.858 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.187), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - ROOT %reshape.3313 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.858), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %slice.924 = s32[512]{0:T(512)} slice(%custom-call.17), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %reshape.3334 = s32[4,128]{1,0:T(4,128)} reshape(%slice.924), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} + %transpose.859 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3334), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} + %gather.191 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.23, %transpose.859), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %transpose.858 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.191), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + ROOT %reshape.3333 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.858), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} } %fused_computation.8 (param_0.26: f32[163840,32], param_1.120: s32[1024]) -> f32[512,32] { %param_0.26 = f32[163840,32]{1,0:T(8,128)} parameter(0) %param_1.120 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.25 = s32[1024]{0:T(1024)} custom-call(%param_1.120), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %slice.904 = s32[512]{0:T(512)} slice(%custom-call.25), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %reshape.3322 = s32[4,128]{1,0:T(4,128)} reshape(%slice.904), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %transpose.865 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3322), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %gather.189 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.26, %transpose.865), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %transpose.864 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.189), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - ROOT %reshape.3321 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %slice.932 = s32[512]{0:T(512)} slice(%custom-call.25), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %reshape.3342 = s32[4,128]{1,0:T(4,128)} reshape(%slice.932), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} + %transpose.865 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3342), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} + %gather.193 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.26, %transpose.865), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %transpose.864 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.193), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + ROOT %reshape.3341 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} } %fused_computation.9 (param_0.29: f32[163840,32], param_1.122: s32[1024]) -> f32[512,32] { %param_0.29 = f32[163840,32]{1,0:T(8,128)} parameter(0) %param_1.122 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.27 = s32[1024]{0:T(1024)} custom-call(%param_1.122), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %slice.906 = s32[512]{0:T(512)} slice(%custom-call.27), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %reshape.3330 = s32[4,128]{1,0:T(4,128)} reshape(%slice.906), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %transpose.871 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3330), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %gather.191 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.29, %transpose.871), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %transpose.870 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.191), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - ROOT %reshape.3329 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.870), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %slice.934 = s32[512]{0:T(512)} slice(%custom-call.27), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %reshape.3350 = s32[4,128]{1,0:T(4,128)} reshape(%slice.934), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} + %transpose.871 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3350), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} + %gather.195 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.29, %transpose.871), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %transpose.870 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.195), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + ROOT %reshape.3349 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.870), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} } %fused_computation.10 (param_0.32: bf16[4096,512], param_1.126: s32[4096]) -> bf16[4096,512] { %param_0.32 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.126 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.31 = s32[4096]{0:T(1024)} custom-call(%param_1.126), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.910 = s32[4096]{0:T(1024)} slice(%custom-call.31), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3338 = s32[4096]{0:T(1024)} reshape(%slice.910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.877 = s32[4096]{0:T(1024)} transpose(%reshape.3338), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.193 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.32, %transpose.877), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.876 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.193), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3337 = bf16[4096,512]{1,0:T(8,128)(2,1)} reshape(%transpose.876), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.938 = s32[4096]{0:T(1024)} slice(%custom-call.31), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3358 = s32[4096]{0:T(1024)} reshape(%slice.938), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.877 = s32[4096]{0:T(1024)} transpose(%reshape.3358), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.197 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.32, %transpose.877), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.876 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.197), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3357 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.876), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.11 (param_0.35: bf16[4096,512], param_1.128: s32[4096]) -> bf16[4096,512] { %param_0.35 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.128 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.33 = s32[4096]{0:T(1024)} custom-call(%param_1.128), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.912 = s32[4096]{0:T(1024)} slice(%custom-call.33), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3346 = s32[4096]{0:T(1024)} reshape(%slice.912), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.883 = s32[4096]{0:T(1024)} transpose(%reshape.3346), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.195 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.35, %transpose.883), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.882 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.195), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3345 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.882), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.940 = s32[4096]{0:T(1024)} slice(%custom-call.33), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3366 = s32[4096]{0:T(1024)} reshape(%slice.940), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.883 = s32[4096]{0:T(1024)} transpose(%reshape.3366), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.199 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.35, %transpose.883), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.882 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.199), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3365 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.882), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.12 (param_0.38: bf16[4096,512], param_1.130: s32[4096]) -> bf16[4096,512] { %param_0.38 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.130 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.35 = s32[4096]{0:T(1024)} custom-call(%param_1.130), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.914 = s32[4096]{0:T(1024)} slice(%custom-call.35), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3354 = s32[4096]{0:T(1024)} reshape(%slice.914), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.889 = s32[4096]{0:T(1024)} transpose(%reshape.3354), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.197 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.38, %transpose.889), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.888 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.197), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3353 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.888), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.942 = s32[4096]{0:T(1024)} slice(%custom-call.35), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3374 = s32[4096]{0:T(1024)} reshape(%slice.942), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.889 = s32[4096]{0:T(1024)} transpose(%reshape.3374), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.201 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.38, %transpose.889), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.888 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.201), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3373 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.888), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.13 (param_0.41: bf16[4096,512], param_1.132: s32[4096]) -> bf16[4096,512] { %param_0.41 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.132 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.37 = s32[4096]{0:T(1024)} custom-call(%param_1.132), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.916 = s32[4096]{0:T(1024)} slice(%custom-call.37), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3362 = s32[4096]{0:T(1024)} reshape(%slice.916), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.895 = s32[4096]{0:T(1024)} transpose(%reshape.3362), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.199 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.41, %transpose.895), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.894 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.199), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3361 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.894), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.944 = s32[4096]{0:T(1024)} slice(%custom-call.37), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3382 = s32[4096]{0:T(1024)} reshape(%slice.944), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.895 = s32[4096]{0:T(1024)} transpose(%reshape.3382), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.203 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.41, %transpose.895), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.894 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.203), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3381 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.894), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.15 (param_0.47: s32[256], param_1.124: s32[1024]) -> s32[263] { %param_0.47 = s32[256]{0:T(256)S(1)} parameter(0) %param_1.124 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.29 = s32[1024]{0:T(1024)} custom-call(%param_1.124), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %slice.908 = s32[263]{0:T(512)} slice(%custom-call.29), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3393 = s32[263]{0:T(512)} reshape(%slice.908), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.911 = s32[263]{0:T(512)} transpose(%reshape.3393), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.204 = s32[263]{0:T(512)} gather(%param_0.47, %transpose.911), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.910 = s32[263]{0:T(512)} transpose(%gather.204), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3392 = s32[263]{0:T(512)S(1)} reshape(%transpose.910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %slice.936 = s32[263]{0:T(512)} slice(%custom-call.29), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.3413 = s32[263]{0:T(512)} reshape(%slice.936), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.911 = s32[263]{0:T(512)} transpose(%reshape.3413), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.208 = s32[263]{0:T(512)} gather(%param_0.47, %transpose.911), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.910 = s32[263]{0:T(512)} transpose(%gather.208), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.3412 = s32[263]{0:T(512)S(1)} reshape(%transpose.910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} } %fused_computation.16 (param_0.50: s32[256], param_1.134: s32[1024]) -> s32[263] { %param_0.50 = s32[256]{0:T(256)S(1)} parameter(0) %param_1.134 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.39 = s32[1024]{0:T(1024)} custom-call(%param_1.134), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %slice.918 = s32[263]{0:T(512)} slice(%custom-call.39), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3416 = s32[263]{0:T(512)} reshape(%slice.918), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.921 = s32[263]{0:T(512)} transpose(%reshape.3416), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.207 = s32[263]{0:T(512)} gather(%param_0.50, %transpose.921), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.920 = s32[263]{0:T(512)} transpose(%gather.207), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3415 = s32[263]{0:T(512)S(1)} reshape(%transpose.920), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %slice.946 = s32[263]{0:T(512)} slice(%custom-call.39), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.3436 = s32[263]{0:T(512)} reshape(%slice.946), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.921 = s32[263]{0:T(512)} transpose(%reshape.3436), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.211 = s32[263]{0:T(512)} gather(%param_0.50, %transpose.921), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.920 = s32[263]{0:T(512)} transpose(%gather.211), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.3435 = s32[263]{0:T(512)S(1)} reshape(%transpose.920), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} } %region_173.198.clone (scatter-add.94: bf16[], scatter-add.96: bf16[]) -> bf16[] { %scatter-add.94 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} %scatter-add.96 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - ROOT %add.1885 = bf16[]{:T(256)} add(%scatter-add.94, %scatter-add.96), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.1875 = bf16[]{:T(256)} add(%scatter-add.94, %scatter-add.96), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %fused_computation.21 (param_0.55: bf16[129280,512], param_1.65: s32[512], param_2.24: bf16[512,512]) -> bf16[129280,512] { %param_0.55 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) %param_1.65 = s32[512]{0:T(512)S(1)} parameter(1) - %reshape.3470 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.65), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.954 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3470), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %reshape.3490 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.65), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %transpose.954 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3490), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} %param_2.24 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} parameter(2) - %reshape.3471 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} reshape(%param_2.24), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} - %transpose.955 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%reshape.3471), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} - ROOT %scatter.73 = bf16[129280,512]{1,0:T(8,128)(2,1)} scatter(%param_0.55, %transpose.954, %transpose.955), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_173.198.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} + %reshape.3491 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} reshape(%param_2.24), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} + %transpose.955 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%reshape.3491), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} + ROOT %scatter.77 = bf16[129280,512]{1,0:T(8,128)(2,1)} scatter(%param_0.55, %transpose.954, %transpose.955), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_173.198.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} } %region_12.18 (top_k.0: bf16[], top_k.6: bf16[], top_k.7: s32[], top_k.8: s32[]) -> pred[] { - %constant.1385 = s32[]{:T(128)} constant(0) - %constant.1386 = s32[]{:T(128)} constant(2147483647) + %constant.1369 = s32[]{:T(128)} constant(0) + %constant.1370 = s32[]{:T(128)} constant(2147483647) %top_k.0 = bf16[]{:T(256)} parameter(0), metadata={op_name="top_k"} %top_k.6 = bf16[]{:T(256)} parameter(1), metadata={op_name="top_k"} %top_k.7 = s32[]{:T(128)} parameter(2), metadata={op_name="top_k"} %top_k.8 = s32[]{:T(128)} parameter(3), metadata={op_name="top_k"} %convert.385 = f32[]{:T(128)S(6)} convert(%top_k.0), metadata={op_name="convert.16"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %bitcast-convert.35 = s32[]{:T(128)S(6)} bitcast-convert(%convert.385), metadata={op_name="bitcast-convert.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.128 = pred[]{:T(512)S(6)} compare(%bitcast-convert.35, %constant.1385), direction=LT, metadata={op_name="compare.35"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.36 = s32[]{:T(128)S(6)} xor(%constant.1386, %bitcast-convert.35), metadata={op_name="xor.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.128 = pred[]{:T(512)S(6)} compare(%bitcast-convert.35, %constant.1369), direction=LT, metadata={op_name="compare.35"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.36 = s32[]{:T(128)S(6)} xor(%constant.1370, %bitcast-convert.35), metadata={op_name="xor.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.118 = s32[]{:T(128)S(6)} select(%compare.128, %xor.36, %bitcast-convert.35), metadata={op_name="select.14"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %convert.386 = f32[]{:T(128)S(6)} convert(%top_k.6), metadata={op_name="convert.17"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %bitcast-convert.36 = s32[]{:T(128)S(6)} bitcast-convert(%convert.386), metadata={op_name="bitcast-convert.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.129 = pred[]{:T(512)S(6)} compare(%bitcast-convert.36, %constant.1385), direction=LT, metadata={op_name="compare.36"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.37 = s32[]{:T(128)S(6)} xor(%constant.1386, %bitcast-convert.36), metadata={op_name="xor.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.129 = pred[]{:T(512)S(6)} compare(%bitcast-convert.36, %constant.1369), direction=LT, metadata={op_name="compare.36"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.37 = s32[]{:T(128)S(6)} xor(%constant.1370, %bitcast-convert.36), metadata={op_name="xor.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.119 = s32[]{:T(128)S(6)} select(%compare.129, %xor.37, %bitcast-convert.36), metadata={op_name="select.15"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %compare.130 = pred[]{:T(512)S(6)} compare(%select.118, %select.119), direction=GT, metadata={op_name="compare.1"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %compare.131 = pred[]{:T(512)S(6)} compare(%select.119, %select.118), direction=GT, metadata={op_name="compare.108"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} @@ -420,12 +420,12 @@ StackFrames %param_0.68 = s32[256]{0:T(256)S(1)} parameter(0) %param_1.114 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.19 = s32[1024]{0:T(1024)} custom-call(%param_1.114), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %slice.898 = s32[263]{0:T(512)} slice(%custom-call.19), slice={[0:263]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3614 = s32[263]{0:T(512)} reshape(%slice.898), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.1037 = s32[263]{0:T(512)} transpose(%reshape.3614), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.209 = s32[263]{0:T(512)} gather(%param_0.68, %transpose.1037), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.1036 = s32[263]{0:T(512)} transpose(%gather.209), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3613 = s32[263]{0:T(512)S(1)} reshape(%transpose.1036), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %slice.926 = s32[263]{0:T(512)} slice(%custom-call.19), slice={[0:263]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.3634 = s32[263]{0:T(512)} reshape(%slice.926), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.1037 = s32[263]{0:T(512)} transpose(%reshape.3634), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.213 = s32[263]{0:T(512)} gather(%param_0.68, %transpose.1037), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.1036 = s32[263]{0:T(512)} transpose(%gather.213), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.3633 = s32[263]{0:T(512)S(1)} reshape(%transpose.1036), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} } %region_27.34.clone.1 (reduce-window.350: s32[], reduce-window.351: s32[]) -> s32[] { @@ -464,12 +464,12 @@ StackFrames %param_0.71 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.116 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.21 = s32[4096]{0:T(1024)} custom-call(%param_1.116), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.900 = s32[4096]{0:T(1024)} slice(%custom-call.21), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3637 = s32[4096]{0:T(1024)} reshape(%slice.900), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.1043 = s32[4096]{0:T(1024)} transpose(%reshape.3637), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.210 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.71, %transpose.1043), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.1042 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.210), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3636 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1042), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.928 = s32[4096]{0:T(1024)} slice(%custom-call.21), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3657 = s32[4096]{0:T(1024)} reshape(%slice.928), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.1043 = s32[4096]{0:T(1024)} transpose(%reshape.3657), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.214 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.71, %transpose.1043), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.1042 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.214), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3656 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1042), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %region_31.39 (sort.50: s32[], sort.51: s32[], sort.52: s32[], sort.53: s32[], sort.54: s32[], sort.55: s32[]) -> pred[] { @@ -490,12 +490,12 @@ StackFrames %param_0.72 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.118 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.23 = s32[4096]{0:T(1024)} custom-call(%param_1.118), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.902 = s32[4096]{0:T(1024)} slice(%custom-call.23), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3639 = s32[4096]{0:T(1024)} reshape(%slice.902), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.1045 = s32[4096]{0:T(1024)} transpose(%reshape.3639), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.211 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.72, %transpose.1045), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.1044 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.211), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3638 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1044), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.930 = s32[4096]{0:T(1024)} slice(%custom-call.23), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3659 = s32[4096]{0:T(1024)} reshape(%slice.930), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.1045 = s32[4096]{0:T(1024)} transpose(%reshape.3659), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.215 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.72, %transpose.1045), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.1044 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.215), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3658 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1044), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %compare (name: s32[], name.1: s32[], name.2: bf16[], name.3: bf16[]) -> pred[] { @@ -503,7 +503,7 @@ StackFrames %name.3 = bf16[] parameter(3) %name = s32[] parameter(0) %name.1 = s32[] parameter(1) - ROOT %compare.377 = pred[] compare(%name, %name.1), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.385 = pred[] compare(%name, %name.1), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %compare.1 (name.4: s32[], name.5: s32[], name.6: f32[], name.7: f32[]) -> pred[] { @@ -511,7 +511,7 @@ StackFrames %name.7 = f32[] parameter(3) %name.4 = s32[] parameter(0) %name.5 = s32[] parameter(1) - ROOT %compare.378 = pred[] compare(%name.4, %name.5), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.386 = pred[] compare(%name.4, %name.5), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %compare.2 (name.8: s32[], name.9: s32[], name.10: f32[], name.11: f32[]) -> pred[] { @@ -519,7 +519,7 @@ StackFrames %name.11 = f32[] parameter(3) %name.8 = s32[] parameter(0) %name.9 = s32[] parameter(1) - ROOT %compare.379 = pred[] compare(%name.8, %name.9), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.387 = pred[] compare(%name.8, %name.9), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %compare.3 (name.12: s32[], name.13: s32[], name.14: f32[], name.15: f32[]) -> pred[] { @@ -527,7 +527,7 @@ StackFrames %name.15 = f32[] parameter(3) %name.12 = s32[] parameter(0) %name.13 = s32[] parameter(1) - ROOT %compare.380 = pred[] compare(%name.12, %name.13), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.388 = pred[] compare(%name.12, %name.13), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %compare.4 (name.16: s32[], name.17: s32[], name.18: f32[], name.19: f32[]) -> pred[] { @@ -535,52 +535,52 @@ StackFrames %name.19 = f32[] parameter(3) %name.16 = s32[] parameter(0) %name.17 = s32[] parameter(1) - ROOT %compare.381 = pred[] compare(%name.16, %name.17), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.389 = pred[] compare(%name.16, %name.17), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%called_computation.13 (param_0.4523: s32[256]) -> s32[256] { - %param_0.4523 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2073 = s32[256]{0:T(256)} copy(%param_0.4523), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"1134","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.13 (param_0.4524: s32[256]) -> s32[256] { + %param_0.4524 = s32[256]{0:T(256)} parameter(0) + ROOT %copy.2073 = s32[256]{0:T(256)} copy(%param_0.4524), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"1134","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.13 (param_0.4524: s32[256]) -> s32[256] { - %param_0.4524 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2074.cloned.1 = s32[256]{0:T(256)} call(%param_0.4524), to_apply=%called_computation.13 +%async_computation.13 (param_0.4525: s32[256]) -> s32[256] { + %param_0.4525 = s32[256]{0:T(256)} parameter(0) + ROOT %copy.2074.cloned.1 = s32[256]{0:T(256)} call(%param_0.4525), to_apply=%called_computation.13 }, execution_thread="sparsecore" %region_49.59 (scatter-add.14: s32[], scatter-add.15: s32[]) -> s32[] { %scatter-add.14 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.15 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1362 = s32[]{:T(128)S(7)} add(%scatter-add.14, %scatter-add.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.22.clone.clone (param_0.4525: s32[256], param_1.5325: s32[4096], param_2.4494: s32[4096]) -> s32[256] { - %param_0.4525 = s32[256]{0:T(256)} parameter(0) - %param_1.5325 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3903 = s32[4096]{0:T(1024)} reshape(%param_1.5325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} - %transpose.1100 = s32[4096]{0:T(1024)} transpose(%reshape.3903), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} - %param_2.4494 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3904 = s32[4096]{0:T(1024)} reshape(%param_2.4494), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - %transpose.1101 = s32[4096]{0:T(1024)} transpose(%reshape.3904), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.231 = s32[256]{0:T(256)} scatter(%param_0.4525, %transpose.1100, %transpose.1101), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_49.59, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %add.1352 = s32[]{:T(128)S(7)} add(%scatter-add.14, %scatter-add.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.14 (param_0.4526: s32[256], param_1.5326: s32[4096], param_2.4495: s32[4096]) -> s32[256] { +%fused_computation.22.clone.clone (param_0.4526: s32[256], param_1.5321: s32[4096], param_2.4492: s32[4096]) -> s32[256] { %param_0.4526 = s32[256]{0:T(256)} parameter(0) - %param_1.5326 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4495 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4526, %param_1.5326, %param_2.4495), kind=kCustom, calls=%fused_computation.22.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5321 = s32[4096]{0:T(1024)} parameter(1) + %reshape.3923 = s32[4096]{0:T(1024)} reshape(%param_1.5321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} + %transpose.1100 = s32[4096]{0:T(1024)} transpose(%reshape.3923), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} + %param_2.4492 = s32[4096]{0:T(1024)} parameter(2) + %reshape.3924 = s32[4096]{0:T(1024)} reshape(%param_2.4492), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + %transpose.1101 = s32[4096]{0:T(1024)} transpose(%reshape.3924), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.237 = s32[256]{0:T(256)} scatter(%param_0.4526, %transpose.1100, %transpose.1101), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_49.59, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.14 (param_0.4527: s32[256], param_1.5327: s32[4096], param_2.4496: s32[4096]) -> s32[256] { +%called_computation.14 (param_0.4527: s32[256], param_1.5322: s32[4096], param_2.4493: s32[4096]) -> s32[256] { %param_0.4527 = s32[256]{0:T(256)} parameter(0) - %param_1.5327 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4496 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4527, %param_1.5327, %param_2.4496), to_apply=%called_computation.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} + %param_1.5322 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4493 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4527, %param_1.5322, %param_2.4493), kind=kCustom, calls=%fused_computation.22.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation (param_0.84: s32[256], param_1.136: s32[4096], param_2.80: s32[4096], param_3.3085: token[]) -> s32[256] { - %param_3.3085 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.14 (param_0.4528: s32[256], param_1.5323: s32[4096], param_2.4494: s32[4096]) -> s32[256] { + %param_0.4528 = s32[256]{0:T(256)} parameter(0) + %param_1.5323 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4494 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4528, %param_1.5323, %param_2.4494), to_apply=%called_computation.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation (param_0.84: s32[256], param_1.136: s32[4096], param_2.80: s32[4096], param_3.3083: token[]) -> s32[256] { + %param_3.3083 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.84 = s32[256]{0:T(256)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.136 = s32[4096]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.80 = s32[4096]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -590,57 +590,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.40.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.40.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation (param_0.85: s32[256], param_1.137: s32[4096], param_2.81: s32[4096], param_3.3084: token[]) -> s32[256] { - %param_3.3084 = token[] parameter(3) +%async_computation (param_0.85: s32[256], param_1.137: s32[4096], param_2.81: s32[4096], param_3.3082: token[]) -> s32[256] { + %param_3.3082 = token[] parameter(3) %param_0.85 = s32[256]{0:T(256)} parameter(0) %param_1.137 = s32[4096]{0:T(1024)} parameter(1) %param_2.81 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.2.cloned.1 = s32[256]{0:T(256)} call(%param_0.85, %param_1.137, %param_2.81, %param_3.3084), to_apply=%called_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.2.cloned.1 = s32[256]{0:T(256)} call(%param_0.85, %param_1.137, %param_2.81, %param_3.3082), to_apply=%called_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.15 (param_0.4528: f32[9]) -> f32[9] { - %param_0.4528 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2075 = f32[9]{0:T(128)} copy(%param_0.4528), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.15 (param_0.4529: f32[9]) -> f32[9] { + %param_0.4529 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2075 = f32[9]{0:T(128)} copy(%param_0.4529), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.15 (param_0.4529: f32[9]) -> f32[9] { - %param_0.4529 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2076.cloned.1 = f32[9]{0:T(128)} call(%param_0.4529), to_apply=%called_computation.15 +%async_computation.15 (param_0.4530: f32[9]) -> f32[9] { + %param_0.4530 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2076.cloned.1 = f32[9]{0:T(128)} call(%param_0.4530), to_apply=%called_computation.15 }, execution_thread="sparsecore" %region_61.72 (scatter-add.24: f32[], scatter-add.25: f32[]) -> f32[] { %scatter-add.24 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.25 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1368 = f32[]{:T(128)S(7)} add(%scatter-add.24, %scatter-add.25), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1358 = f32[]{:T(128)S(7)} add(%scatter-add.24, %scatter-add.25), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.24.clone.clone (param_0.4530: f32[9], param_1.5328: s32[256], param_2.4497: f32[256]) -> f32[9] { - %param_0.4530 = f32[9]{0:T(128)} parameter(0) - %param_1.5328 = s32[256]{0:T(256)} parameter(1) - %reshape.3905 = s32[256]{0:T(256)} reshape(%param_1.5328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1102 = s32[256]{0:T(256)} transpose(%reshape.3905), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4497 = f32[256]{0:T(256)} parameter(2) - %reshape.3906 = f32[256]{0:T(256)} reshape(%param_2.4497), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1103 = f32[256]{0:T(256)} transpose(%reshape.3906), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.232 = f32[9]{0:T(128)} scatter(%param_0.4530, %transpose.1102, %transpose.1103), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_61.72, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.16 (param_0.4531: f32[9], param_1.5329: s32[256], param_2.4498: f32[256]) -> f32[9] { +%fused_computation.24.clone.clone (param_0.4531: f32[9], param_1.5324: s32[256], param_2.4495: f32[256]) -> f32[9] { %param_0.4531 = f32[9]{0:T(128)} parameter(0) - %param_1.5329 = s32[256]{0:T(256)} parameter(1) - %param_2.4498 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.41 = f32[9]{0:T(128)} fusion(%param_0.4531, %param_1.5329, %param_2.4498), kind=kCustom, calls=%fused_computation.24.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5324 = s32[256]{0:T(256)} parameter(1) + %reshape.3925 = s32[256]{0:T(256)} reshape(%param_1.5324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1102 = s32[256]{0:T(256)} transpose(%reshape.3925), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4495 = f32[256]{0:T(256)} parameter(2) + %reshape.3926 = f32[256]{0:T(256)} reshape(%param_2.4495), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1103 = f32[256]{0:T(256)} transpose(%reshape.3926), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.238 = f32[9]{0:T(128)} scatter(%param_0.4531, %transpose.1102, %transpose.1103), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_61.72, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.16 (param_0.4532: f32[9], param_1.5330: s32[256], param_2.4499: f32[256]) -> f32[9] { +%called_computation.16 (param_0.4532: f32[9], param_1.5325: s32[256], param_2.4496: f32[256]) -> f32[9] { %param_0.4532 = f32[9]{0:T(128)} parameter(0) - %param_1.5330 = s32[256]{0:T(256)} parameter(1) - %param_2.4499 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.42.cloned.1 = f32[9]{0:T(128)} call(%param_0.4532, %param_1.5330, %param_2.4499), to_apply=%called_computation.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5325 = s32[256]{0:T(256)} parameter(1) + %param_2.4496 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.41 = f32[9]{0:T(128)} fusion(%param_0.4532, %param_1.5325, %param_2.4496), kind=kCustom, calls=%fused_computation.24.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.1 (param_0.87: f32[9], param_1.139: s32[256], param_2.83: f32[256], param_3.3099: token[]) -> f32[9] { - %param_3.3099 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.16 (param_0.4533: f32[9], param_1.5326: s32[256], param_2.4497: f32[256]) -> f32[9] { + %param_0.4533 = f32[9]{0:T(128)} parameter(0) + %param_1.5326 = s32[256]{0:T(256)} parameter(1) + %param_2.4497 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.42.cloned.1 = f32[9]{0:T(128)} call(%param_0.4533, %param_1.5326, %param_2.4497), to_apply=%called_computation.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.1 (param_0.87: f32[9], param_1.139: s32[256], param_2.83: f32[256], param_3.3097: token[]) -> f32[9] { + %param_3.3097 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.87 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.139 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.83 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -650,57 +650,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.42.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.42.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.1 (param_0.88: f32[9], param_1.140: s32[256], param_2.84: f32[256], param_3.3098: token[]) -> f32[9] { - %param_3.3098 = token[] parameter(3) +%async_computation.1 (param_0.88: f32[9], param_1.140: s32[256], param_2.84: f32[256], param_3.3096: token[]) -> f32[9] { + %param_3.3096 = token[] parameter(3) %param_0.88 = f32[9]{0:T(128)} parameter(0) %param_1.140 = s32[256]{0:T(256)} parameter(1) %param_2.84 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.5.cloned.1 = f32[9]{0:T(128)} call(%param_0.88, %param_1.140, %param_2.84, %param_3.3098), to_apply=%called_computation.1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.5.cloned.1 = f32[9]{0:T(128)} call(%param_0.88, %param_1.140, %param_2.84, %param_3.3096), to_apply=%called_computation.1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.17 (param_0.4533: s32[263]) -> s32[263] { - %param_0.4533 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2077 = s32[263]{0:T(512)} copy(%param_0.4533), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.17 (param_0.4534: s32[263]) -> s32[263] { + %param_0.4534 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2077 = s32[263]{0:T(512)} copy(%param_0.4534), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.17 (param_0.4534: s32[263]) -> s32[263] { - %param_0.4534 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2078.cloned.1 = s32[263]{0:T(512)} call(%param_0.4534), to_apply=%called_computation.17 +%async_computation.17 (param_0.4535: s32[263]) -> s32[263] { + %param_0.4535 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2078.cloned.1 = s32[263]{0:T(512)} call(%param_0.4535), to_apply=%called_computation.17 }, execution_thread="sparsecore" %region_63.74 (scatter-add.28: s32[], scatter-add.29: s32[]) -> s32[] { %scatter-add.28 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.29 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1369 = s32[]{:T(128)S(7)} add(%scatter-add.28, %scatter-add.29), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1359 = s32[]{:T(128)S(7)} add(%scatter-add.28, %scatter-add.29), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.25.clone.clone (param_0.4535: s32[263], param_1.5331: s32[8], param_2.4500: s32[8]) -> s32[263] { - %param_0.4535 = s32[263]{0:T(512)} parameter(0) - %param_1.5331 = s32[8]{0:T(128)} parameter(1) - %reshape.3907 = s32[8]{0:T(128)} reshape(%param_1.5331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1104 = s32[8]{0:T(128)} transpose(%reshape.3907), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4500 = s32[8]{0:T(128)} parameter(2) - %reshape.3908 = s32[8]{0:T(128)} reshape(%param_2.4500), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1105 = s32[8]{0:T(128)} transpose(%reshape.3908), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.233 = s32[263]{0:T(512)} scatter(%param_0.4535, %transpose.1104, %transpose.1105), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_63.74, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.18 (param_0.4536: s32[263], param_1.5332: s32[8], param_2.4501: s32[8]) -> s32[263] { +%fused_computation.25.clone.clone (param_0.4536: s32[263], param_1.5327: s32[8], param_2.4498: s32[8]) -> s32[263] { %param_0.4536 = s32[263]{0:T(512)} parameter(0) - %param_1.5332 = s32[8]{0:T(128)} parameter(1) - %param_2.4501 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.43 = s32[263]{0:T(512)} fusion(%param_0.4536, %param_1.5332, %param_2.4501), kind=kCustom, calls=%fused_computation.25.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5327 = s32[8]{0:T(128)} parameter(1) + %reshape.3927 = s32[8]{0:T(128)} reshape(%param_1.5327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1104 = s32[8]{0:T(128)} transpose(%reshape.3927), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4498 = s32[8]{0:T(128)} parameter(2) + %reshape.3928 = s32[8]{0:T(128)} reshape(%param_2.4498), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %transpose.1105 = s32[8]{0:T(128)} transpose(%reshape.3928), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + ROOT %scatter-add.239 = s32[263]{0:T(512)} scatter(%param_0.4536, %transpose.1104, %transpose.1105), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_63.74, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.18 (param_0.4537: s32[263], param_1.5333: s32[8], param_2.4502: s32[8]) -> s32[263] { +%called_computation.18 (param_0.4537: s32[263], param_1.5328: s32[8], param_2.4499: s32[8]) -> s32[263] { %param_0.4537 = s32[263]{0:T(512)} parameter(0) - %param_1.5333 = s32[8]{0:T(128)} parameter(1) - %param_2.4502 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.44.cloned.1 = s32[263]{0:T(512)} call(%param_0.4537, %param_1.5333, %param_2.4502), to_apply=%called_computation.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5328 = s32[8]{0:T(128)} parameter(1) + %param_2.4499 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.43 = s32[263]{0:T(512)} fusion(%param_0.4537, %param_1.5328, %param_2.4499), kind=kCustom, calls=%fused_computation.25.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.2 (param_0.90: s32[263], param_1.142: s32[8], param_2.86: s32[8], param_3.3105: token[]) -> s32[263] { - %param_3.3105 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.18 (param_0.4538: s32[263], param_1.5329: s32[8], param_2.4500: s32[8]) -> s32[263] { + %param_0.4538 = s32[263]{0:T(512)} parameter(0) + %param_1.5329 = s32[8]{0:T(128)} parameter(1) + %param_2.4500 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.44.cloned.1 = s32[263]{0:T(512)} call(%param_0.4538, %param_1.5329, %param_2.4500), to_apply=%called_computation.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.2 (param_0.90: s32[263], param_1.142: s32[8], param_2.86: s32[8], param_3.3103: token[]) -> s32[263] { + %param_3.3103 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.90 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.142 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.86 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -710,57 +710,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.44.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.44.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.2 (param_0.91: s32[263], param_1.143: s32[8], param_2.87: s32[8], param_3.3104: token[]) -> s32[263] { - %param_3.3104 = token[] parameter(3) +%async_computation.2 (param_0.91: s32[263], param_1.143: s32[8], param_2.87: s32[8], param_3.3102: token[]) -> s32[263] { + %param_3.3102 = token[] parameter(3) %param_0.91 = s32[263]{0:T(512)} parameter(0) %param_1.143 = s32[8]{0:T(128)} parameter(1) %param_2.87 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.8.cloned.1 = s32[263]{0:T(512)} call(%param_0.91, %param_1.143, %param_2.87, %param_3.3104), to_apply=%called_computation.2, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.8.cloned.1 = s32[263]{0:T(512)} call(%param_0.91, %param_1.143, %param_2.87, %param_3.3102), to_apply=%called_computation.2, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.19 (param_0.4538: s32[263]) -> s32[263] { - %param_0.4538 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2079 = s32[263]{0:T(512)} copy(%param_0.4538), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.19 (param_0.4539: s32[263]) -> s32[263] { + %param_0.4539 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2079 = s32[263]{0:T(512)} copy(%param_0.4539), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.19 (param_0.4539: s32[263]) -> s32[263] { - %param_0.4539 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2080.cloned.1 = s32[263]{0:T(512)} call(%param_0.4539), to_apply=%called_computation.19 +%async_computation.19 (param_0.4540: s32[263]) -> s32[263] { + %param_0.4540 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2080.cloned.1 = s32[263]{0:T(512)} call(%param_0.4540), to_apply=%called_computation.19 }, execution_thread="sparsecore" %region_73.86.clone (scatter-add.163: s32[], scatter-add.164: s32[]) -> s32[] { %scatter-add.163 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.164 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2482 = s32[]{:T(128)S(7)} add(%scatter-add.163, %scatter-add.164), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.26.clone.clone (param_0.4540: s32[263], param_1.5334: s32[256], param_2.4503: s32[256]) -> s32[263] { - %param_0.4540 = s32[263]{0:T(512)} parameter(0) - %param_1.5334 = s32[256]{0:T(256)} parameter(1) - %reshape.3909 = s32[256]{0:T(256)} reshape(%param_1.5334), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1106 = s32[256]{0:T(256)} transpose(%reshape.3909), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4503 = s32[256]{0:T(256)} parameter(2) - %reshape.3910 = s32[256]{0:T(256)} reshape(%param_2.4503), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1107 = s32[256]{0:T(256)} transpose(%reshape.3910), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.234 = s32[263]{0:T(512)} scatter(%param_0.4540, %transpose.1106, %transpose.1107), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_73.86.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %add.2474 = s32[]{:T(128)S(7)} add(%scatter-add.163, %scatter-add.164), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.20 (param_0.4541: s32[263], param_1.5335: s32[256], param_2.4504: s32[256]) -> s32[263] { +%fused_computation.26.clone.clone (param_0.4541: s32[263], param_1.5330: s32[256], param_2.4501: s32[256]) -> s32[263] { %param_0.4541 = s32[263]{0:T(512)} parameter(0) - %param_1.5335 = s32[256]{0:T(256)} parameter(1) - %param_2.4504 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.45 = s32[263]{0:T(512)} fusion(%param_0.4541, %param_1.5335, %param_2.4504), kind=kCustom, calls=%fused_computation.26.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5330 = s32[256]{0:T(256)} parameter(1) + %reshape.3929 = s32[256]{0:T(256)} reshape(%param_1.5330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1106 = s32[256]{0:T(256)} transpose(%reshape.3929), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4501 = s32[256]{0:T(256)} parameter(2) + %reshape.3930 = s32[256]{0:T(256)} reshape(%param_2.4501), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1107 = s32[256]{0:T(256)} transpose(%reshape.3930), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.240 = s32[263]{0:T(512)} scatter(%param_0.4541, %transpose.1106, %transpose.1107), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_73.86.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.20 (param_0.4542: s32[263], param_1.5336: s32[256], param_2.4505: s32[256]) -> s32[263] { +%called_computation.20 (param_0.4542: s32[263], param_1.5331: s32[256], param_2.4502: s32[256]) -> s32[263] { %param_0.4542 = s32[263]{0:T(512)} parameter(0) - %param_1.5336 = s32[256]{0:T(256)} parameter(1) - %param_2.4505 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.46.cloned.1 = s32[263]{0:T(512)} call(%param_0.4542, %param_1.5336, %param_2.4505), to_apply=%called_computation.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5331 = s32[256]{0:T(256)} parameter(1) + %param_2.4502 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.45 = s32[263]{0:T(512)} fusion(%param_0.4542, %param_1.5331, %param_2.4502), kind=kCustom, calls=%fused_computation.26.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.3 (param_0.93: s32[263], param_1.145: s32[256], param_2.89: s32[256], param_3.3091: token[]) -> s32[263] { - %param_3.3091 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.20 (param_0.4543: s32[263], param_1.5332: s32[256], param_2.4503: s32[256]) -> s32[263] { + %param_0.4543 = s32[263]{0:T(512)} parameter(0) + %param_1.5332 = s32[256]{0:T(256)} parameter(1) + %param_2.4503 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.46.cloned.1 = s32[263]{0:T(512)} call(%param_0.4543, %param_1.5332, %param_2.4503), to_apply=%called_computation.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.3 (param_0.93: s32[263], param_1.145: s32[256], param_2.89: s32[256], param_3.3089: token[]) -> s32[263] { + %param_3.3089 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.93 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.145 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.89 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -770,57 +770,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.46.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.46.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.3 (param_0.94: s32[263], param_1.146: s32[256], param_2.90: s32[256], param_3.3090: token[]) -> s32[263] { - %param_3.3090 = token[] parameter(3) +%async_computation.3 (param_0.94: s32[263], param_1.146: s32[256], param_2.90: s32[256], param_3.3088: token[]) -> s32[263] { + %param_3.3088 = token[] parameter(3) %param_0.94 = s32[263]{0:T(512)} parameter(0) %param_1.146 = s32[256]{0:T(256)} parameter(1) %param_2.90 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.11.cloned.1 = s32[263]{0:T(512)} call(%param_0.94, %param_1.146, %param_2.90, %param_3.3090), to_apply=%called_computation.3, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.11.cloned.1 = s32[263]{0:T(512)} call(%param_0.94, %param_1.146, %param_2.90, %param_3.3088), to_apply=%called_computation.3, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.21 (param_0.4543: f32[9]) -> f32[9] { - %param_0.4543 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2081 = f32[9]{0:T(128)} copy(%param_0.4543), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.21 (param_0.4544: f32[9]) -> f32[9] { + %param_0.4544 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2081 = f32[9]{0:T(128)} copy(%param_0.4544), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.21 (param_0.4544: f32[9]) -> f32[9] { - %param_0.4544 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2082.cloned.1 = f32[9]{0:T(128)} call(%param_0.4544), to_apply=%called_computation.21 +%async_computation.21 (param_0.4545: f32[9]) -> f32[9] { + %param_0.4545 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2082.cloned.1 = f32[9]{0:T(128)} call(%param_0.4545), to_apply=%called_computation.21 }, execution_thread="sparsecore" %region_79.95.clone (scatter-add.167: f32[], scatter-add.168: f32[]) -> f32[] { %scatter-add.167 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.168 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2484 = f32[]{:T(128)S(7)} add(%scatter-add.167, %scatter-add.168), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.2476 = f32[]{:T(128)S(7)} add(%scatter-add.167, %scatter-add.168), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.27.clone.clone (param_0.4545: f32[9], param_1.5337: s32[256], param_2.4506: f32[256]) -> f32[9] { - %param_0.4545 = f32[9]{0:T(128)} parameter(0) - %param_1.5337 = s32[256]{0:T(256)} parameter(1) - %reshape.3911 = s32[256]{0:T(256)} reshape(%param_1.5337), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1108 = s32[256]{0:T(256)} transpose(%reshape.3911), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4506 = f32[256]{0:T(256)} parameter(2) - %reshape.3912 = f32[256]{0:T(256)} reshape(%param_2.4506), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1109 = f32[256]{0:T(256)} transpose(%reshape.3912), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.235 = f32[9]{0:T(128)} scatter(%param_0.4545, %transpose.1108, %transpose.1109), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_79.95.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.22 (param_0.4546: f32[9], param_1.5338: s32[256], param_2.4507: f32[256]) -> f32[9] { +%fused_computation.27.clone.clone (param_0.4546: f32[9], param_1.5333: s32[256], param_2.4504: f32[256]) -> f32[9] { %param_0.4546 = f32[9]{0:T(128)} parameter(0) - %param_1.5338 = s32[256]{0:T(256)} parameter(1) - %param_2.4507 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.47 = f32[9]{0:T(128)} fusion(%param_0.4546, %param_1.5338, %param_2.4507), kind=kCustom, calls=%fused_computation.27.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5333 = s32[256]{0:T(256)} parameter(1) + %reshape.3931 = s32[256]{0:T(256)} reshape(%param_1.5333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1108 = s32[256]{0:T(256)} transpose(%reshape.3931), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4504 = f32[256]{0:T(256)} parameter(2) + %reshape.3932 = f32[256]{0:T(256)} reshape(%param_2.4504), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1109 = f32[256]{0:T(256)} transpose(%reshape.3932), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.241 = f32[9]{0:T(128)} scatter(%param_0.4546, %transpose.1108, %transpose.1109), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_79.95.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.22 (param_0.4547: f32[9], param_1.5339: s32[256], param_2.4508: f32[256]) -> f32[9] { +%called_computation.22 (param_0.4547: f32[9], param_1.5334: s32[256], param_2.4505: f32[256]) -> f32[9] { %param_0.4547 = f32[9]{0:T(128)} parameter(0) - %param_1.5339 = s32[256]{0:T(256)} parameter(1) - %param_2.4508 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.48.cloned.1 = f32[9]{0:T(128)} call(%param_0.4547, %param_1.5339, %param_2.4508), to_apply=%called_computation.22, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5334 = s32[256]{0:T(256)} parameter(1) + %param_2.4505 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.47 = f32[9]{0:T(128)} fusion(%param_0.4547, %param_1.5334, %param_2.4505), kind=kCustom, calls=%fused_computation.27.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.4 (param_0.96: f32[9], param_1.148: s32[256], param_2.92: f32[256], param_3.3097: token[]) -> f32[9] { - %param_3.3097 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.22 (param_0.4548: f32[9], param_1.5335: s32[256], param_2.4506: f32[256]) -> f32[9] { + %param_0.4548 = f32[9]{0:T(128)} parameter(0) + %param_1.5335 = s32[256]{0:T(256)} parameter(1) + %param_2.4506 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.48.cloned.1 = f32[9]{0:T(128)} call(%param_0.4548, %param_1.5335, %param_2.4506), to_apply=%called_computation.22, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.4 (param_0.96: f32[9], param_1.148: s32[256], param_2.92: f32[256], param_3.3095: token[]) -> f32[9] { + %param_3.3095 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.96 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.148 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.92 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -830,57 +830,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.48.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.48.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.4 (param_0.97: f32[9], param_1.149: s32[256], param_2.93: f32[256], param_3.3096: token[]) -> f32[9] { - %param_3.3096 = token[] parameter(3) +%async_computation.4 (param_0.97: f32[9], param_1.149: s32[256], param_2.93: f32[256], param_3.3094: token[]) -> f32[9] { + %param_3.3094 = token[] parameter(3) %param_0.97 = f32[9]{0:T(128)} parameter(0) %param_1.149 = s32[256]{0:T(256)} parameter(1) %param_2.93 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.14.cloned.1 = f32[9]{0:T(128)} call(%param_0.97, %param_1.149, %param_2.93, %param_3.3096), to_apply=%called_computation.4, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.14.cloned.1 = f32[9]{0:T(128)} call(%param_0.97, %param_1.149, %param_2.93, %param_3.3094), to_apply=%called_computation.4, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.23 (param_0.4548: s32[263]) -> s32[263] { - %param_0.4548 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2083 = s32[263]{0:T(512)} copy(%param_0.4548), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.23 (param_0.4549: s32[263]) -> s32[263] { + %param_0.4549 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2083 = s32[263]{0:T(512)} copy(%param_0.4549), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.23 (param_0.4549: s32[263]) -> s32[263] { - %param_0.4549 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2084.cloned.1 = s32[263]{0:T(512)} call(%param_0.4549), to_apply=%called_computation.23 +%async_computation.23 (param_0.4550: s32[263]) -> s32[263] { + %param_0.4550 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2084.cloned.1 = s32[263]{0:T(512)} call(%param_0.4550), to_apply=%called_computation.23 }, execution_thread="sparsecore" %region_81.97.clone (scatter-add.171: s32[], scatter-add.172: s32[]) -> s32[] { %scatter-add.171 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.172 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2486 = s32[]{:T(128)S(7)} add(%scatter-add.171, %scatter-add.172), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.28.clone.clone (param_0.4550: s32[263], param_1.5340: s32[8], param_2.4509: s32[8]) -> s32[263] { - %param_0.4550 = s32[263]{0:T(512)} parameter(0) - %param_1.5340 = s32[8]{0:T(128)} parameter(1) - %reshape.3913 = s32[8]{0:T(128)} reshape(%param_1.5340), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1110 = s32[8]{0:T(128)} transpose(%reshape.3913), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4509 = s32[8]{0:T(128)} parameter(2) - %reshape.3914 = s32[8]{0:T(128)} reshape(%param_2.4509), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1111 = s32[8]{0:T(128)} transpose(%reshape.3914), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.236 = s32[263]{0:T(512)} scatter(%param_0.4550, %transpose.1110, %transpose.1111), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_81.97.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %add.2478 = s32[]{:T(128)S(7)} add(%scatter-add.171, %scatter-add.172), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.24 (param_0.4551: s32[263], param_1.5341: s32[8], param_2.4510: s32[8]) -> s32[263] { +%fused_computation.28.clone.clone (param_0.4551: s32[263], param_1.5336: s32[8], param_2.4507: s32[8]) -> s32[263] { %param_0.4551 = s32[263]{0:T(512)} parameter(0) - %param_1.5341 = s32[8]{0:T(128)} parameter(1) - %param_2.4510 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.49 = s32[263]{0:T(512)} fusion(%param_0.4551, %param_1.5341, %param_2.4510), kind=kCustom, calls=%fused_computation.28.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5336 = s32[8]{0:T(128)} parameter(1) + %reshape.3933 = s32[8]{0:T(128)} reshape(%param_1.5336), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1110 = s32[8]{0:T(128)} transpose(%reshape.3933), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4507 = s32[8]{0:T(128)} parameter(2) + %reshape.3934 = s32[8]{0:T(128)} reshape(%param_2.4507), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %transpose.1111 = s32[8]{0:T(128)} transpose(%reshape.3934), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + ROOT %scatter-add.242 = s32[263]{0:T(512)} scatter(%param_0.4551, %transpose.1110, %transpose.1111), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_81.97.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.24 (param_0.4552: s32[263], param_1.5342: s32[8], param_2.4511: s32[8]) -> s32[263] { +%called_computation.24 (param_0.4552: s32[263], param_1.5337: s32[8], param_2.4508: s32[8]) -> s32[263] { %param_0.4552 = s32[263]{0:T(512)} parameter(0) - %param_1.5342 = s32[8]{0:T(128)} parameter(1) - %param_2.4511 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.50.cloned.1 = s32[263]{0:T(512)} call(%param_0.4552, %param_1.5342, %param_2.4511), to_apply=%called_computation.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5337 = s32[8]{0:T(128)} parameter(1) + %param_2.4508 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.49 = s32[263]{0:T(512)} fusion(%param_0.4552, %param_1.5337, %param_2.4508), kind=kCustom, calls=%fused_computation.28.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.5 (param_0.99: s32[263], param_1.151: s32[8], param_2.95: s32[8], param_3.3107: token[]) -> s32[263] { - %param_3.3107 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.24 (param_0.4553: s32[263], param_1.5338: s32[8], param_2.4509: s32[8]) -> s32[263] { + %param_0.4553 = s32[263]{0:T(512)} parameter(0) + %param_1.5338 = s32[8]{0:T(128)} parameter(1) + %param_2.4509 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.50.cloned.1 = s32[263]{0:T(512)} call(%param_0.4553, %param_1.5338, %param_2.4509), to_apply=%called_computation.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.5 (param_0.99: s32[263], param_1.151: s32[8], param_2.95: s32[8], param_3.3105: token[]) -> s32[263] { + %param_3.3105 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.99 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.151 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.95 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -890,57 +890,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.50.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.50.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.5 (param_0.100: s32[263], param_1.152: s32[8], param_2.96: s32[8], param_3.3106: token[]) -> s32[263] { - %param_3.3106 = token[] parameter(3) +%async_computation.5 (param_0.100: s32[263], param_1.152: s32[8], param_2.96: s32[8], param_3.3104: token[]) -> s32[263] { + %param_3.3104 = token[] parameter(3) %param_0.100 = s32[263]{0:T(512)} parameter(0) %param_1.152 = s32[8]{0:T(128)} parameter(1) %param_2.96 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.17.cloned.1 = s32[263]{0:T(512)} call(%param_0.100, %param_1.152, %param_2.96, %param_3.3106), to_apply=%called_computation.5, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.17.cloned.1 = s32[263]{0:T(512)} call(%param_0.100, %param_1.152, %param_2.96, %param_3.3104), to_apply=%called_computation.5, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.25 (param_0.4553: s32[263]) -> s32[263] { - %param_0.4553 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2085 = s32[263]{0:T(512)} copy(%param_0.4553), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.25 (param_0.4554: s32[263]) -> s32[263] { + %param_0.4554 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2085 = s32[263]{0:T(512)} copy(%param_0.4554), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.25 (param_0.4554: s32[263]) -> s32[263] { - %param_0.4554 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2086.cloned.1 = s32[263]{0:T(512)} call(%param_0.4554), to_apply=%called_computation.25 +%async_computation.25 (param_0.4555: s32[263]) -> s32[263] { + %param_0.4555 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2086.cloned.1 = s32[263]{0:T(512)} call(%param_0.4555), to_apply=%called_computation.25 }, execution_thread="sparsecore" %region_96.114 (scatter-add.48: s32[], scatter-add.49: s32[]) -> s32[] { %scatter-add.48 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.49 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1406 = s32[]{:T(128)S(7)} add(%scatter-add.48, %scatter-add.49), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1396 = s32[]{:T(128)S(7)} add(%scatter-add.48, %scatter-add.49), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.29.clone.clone (param_0.4555: s32[263], param_1.5343: s32[256], param_2.4512: s32[256]) -> s32[263] { - %param_0.4555 = s32[263]{0:T(512)} parameter(0) - %param_1.5343 = s32[256]{0:T(256)} parameter(1) - %reshape.3915 = s32[256]{0:T(256)} reshape(%param_1.5343), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %transpose.1112 = s32[256]{0:T(256)} transpose(%reshape.3915), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %param_2.4512 = s32[256]{0:T(256)} parameter(2) - %reshape.3916 = s32[256]{0:T(256)} reshape(%param_2.4512), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1113 = s32[256]{0:T(256)} transpose(%reshape.3916), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.237 = s32[263]{0:T(512)} scatter(%param_0.4555, %transpose.1112, %transpose.1113), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_96.114, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.26 (param_0.4556: s32[263], param_1.5344: s32[256], param_2.4513: s32[256]) -> s32[263] { +%fused_computation.29.clone.clone (param_0.4556: s32[263], param_1.5339: s32[256], param_2.4510: s32[256]) -> s32[263] { %param_0.4556 = s32[263]{0:T(512)} parameter(0) - %param_1.5344 = s32[256]{0:T(256)} parameter(1) - %param_2.4513 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.51 = s32[263]{0:T(512)} fusion(%param_0.4556, %param_1.5344, %param_2.4513), kind=kCustom, calls=%fused_computation.29.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5339 = s32[256]{0:T(256)} parameter(1) + %reshape.3935 = s32[256]{0:T(256)} reshape(%param_1.5339), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %transpose.1112 = s32[256]{0:T(256)} transpose(%reshape.3935), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %param_2.4510 = s32[256]{0:T(256)} parameter(2) + %reshape.3936 = s32[256]{0:T(256)} reshape(%param_2.4510), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1113 = s32[256]{0:T(256)} transpose(%reshape.3936), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.243 = s32[263]{0:T(512)} scatter(%param_0.4556, %transpose.1112, %transpose.1113), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_96.114, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.26 (param_0.4557: s32[263], param_1.5345: s32[256], param_2.4514: s32[256]) -> s32[263] { +%called_computation.26 (param_0.4557: s32[263], param_1.5340: s32[256], param_2.4511: s32[256]) -> s32[263] { %param_0.4557 = s32[263]{0:T(512)} parameter(0) - %param_1.5345 = s32[256]{0:T(256)} parameter(1) - %param_2.4514 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.52.cloned.1 = s32[263]{0:T(512)} call(%param_0.4557, %param_1.5345, %param_2.4514), to_apply=%called_computation.26, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + %param_1.5340 = s32[256]{0:T(256)} parameter(1) + %param_2.4511 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.51 = s32[263]{0:T(512)} fusion(%param_0.4557, %param_1.5340, %param_2.4511), kind=kCustom, calls=%fused_computation.29.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.6 (param_0.102: s32[263], param_1.154: s32[256], param_2.98: s32[256], param_3.3093: token[]) -> s32[263] { - %param_3.3093 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.26 (param_0.4558: s32[263], param_1.5341: s32[256], param_2.4512: s32[256]) -> s32[263] { + %param_0.4558 = s32[263]{0:T(512)} parameter(0) + %param_1.5341 = s32[256]{0:T(256)} parameter(1) + %param_2.4512 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.52.cloned.1 = s32[263]{0:T(512)} call(%param_0.4558, %param_1.5341, %param_2.4512), to_apply=%called_computation.26, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.6 (param_0.102: s32[263], param_1.154: s32[256], param_2.98: s32[256], param_3.3091: token[]) -> s32[263] { + %param_3.3091 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.102 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.154 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.98 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -950,47 +950,47 @@ StackFrames ROOT %scatter_offload_custom_fusion.52.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.52.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.6 (param_0.103: s32[263], param_1.155: s32[256], param_2.99: s32[256], param_3.3092: token[]) -> s32[263] { - %param_3.3092 = token[] parameter(3) +%async_computation.6 (param_0.103: s32[263], param_1.155: s32[256], param_2.99: s32[256], param_3.3090: token[]) -> s32[263] { + %param_3.3090 = token[] parameter(3) %param_0.103 = s32[263]{0:T(512)} parameter(0) %param_1.155 = s32[256]{0:T(256)} parameter(1) %param_2.99 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.20.cloned.1 = s32[263]{0:T(512)} call(%param_0.103, %param_1.155, %param_2.99, %param_3.3092), to_apply=%called_computation.6, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.20.cloned.1 = s32[263]{0:T(512)} call(%param_0.103, %param_1.155, %param_2.99, %param_3.3090), to_apply=%called_computation.6, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_102.120 (scatter-add.52: f32[], scatter-add.53: f32[]) -> f32[] { %scatter-add.52 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.53 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1409 = f32[]{:T(128)S(7)} add(%scatter-add.52, %scatter-add.53), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.30.clone.clone (param_0.4560: f32[9], param_1.5346: s32[256], param_2.4515: f32[256]) -> f32[9] { - %param_0.4560 = f32[9]{0:T(128)} parameter(0) - %param_1.5346 = s32[256]{0:T(256)} parameter(1) - %reshape.3917 = s32[256]{0:T(256)} reshape(%param_1.5346), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1114 = s32[256]{0:T(256)} transpose(%reshape.3917), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4515 = f32[256]{0:T(256)} parameter(2) - %reshape.3918 = f32[256]{0:T(256)} reshape(%param_2.4515), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1115 = f32[256]{0:T(256)} transpose(%reshape.3918), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.238 = f32[9]{0:T(128)} scatter(%param_0.4560, %transpose.1114, %transpose.1115), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_102.120, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %add.1399 = f32[]{:T(128)S(7)} add(%scatter-add.52, %scatter-add.53), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.28 (param_0.4561: f32[9], param_1.5347: s32[256], param_2.4516: f32[256]) -> f32[9] { +%fused_computation.30.clone.clone (param_0.4561: f32[9], param_1.5342: s32[256], param_2.4513: f32[256]) -> f32[9] { %param_0.4561 = f32[9]{0:T(128)} parameter(0) - %param_1.5347 = s32[256]{0:T(256)} parameter(1) - %param_2.4516 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.53 = f32[9]{0:T(128)} fusion(%param_0.4561, %param_1.5347, %param_2.4516), kind=kCustom, calls=%fused_computation.30.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5342 = s32[256]{0:T(256)} parameter(1) + %reshape.3937 = s32[256]{0:T(256)} reshape(%param_1.5342), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1114 = s32[256]{0:T(256)} transpose(%reshape.3937), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4513 = f32[256]{0:T(256)} parameter(2) + %reshape.3938 = f32[256]{0:T(256)} reshape(%param_2.4513), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1115 = f32[256]{0:T(256)} transpose(%reshape.3938), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.244 = f32[9]{0:T(128)} scatter(%param_0.4561, %transpose.1114, %transpose.1115), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_102.120, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.28 (param_0.4562: f32[9], param_1.5348: s32[256], param_2.4517: f32[256]) -> f32[9] { +%called_computation.28 (param_0.4562: f32[9], param_1.5343: s32[256], param_2.4514: f32[256]) -> f32[9] { %param_0.4562 = f32[9]{0:T(128)} parameter(0) - %param_1.5348 = s32[256]{0:T(256)} parameter(1) - %param_2.4517 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.54.cloned.1 = f32[9]{0:T(128)} call(%param_0.4562, %param_1.5348, %param_2.4517), to_apply=%called_computation.28, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + %param_1.5343 = s32[256]{0:T(256)} parameter(1) + %param_2.4514 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.53 = f32[9]{0:T(128)} fusion(%param_0.4562, %param_1.5343, %param_2.4514), kind=kCustom, calls=%fused_computation.30.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.7 (param_0.105: f32[9], param_1.157: s32[256], param_2.101: f32[256], param_3.3101: token[]) -> f32[9] { - %param_3.3101 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.28 (param_0.4563: f32[9], param_1.5344: s32[256], param_2.4515: f32[256]) -> f32[9] { + %param_0.4563 = f32[9]{0:T(128)} parameter(0) + %param_1.5344 = s32[256]{0:T(256)} parameter(1) + %param_2.4515 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.54.cloned.1 = f32[9]{0:T(128)} call(%param_0.4563, %param_1.5344, %param_2.4515), to_apply=%called_computation.28, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.7 (param_0.105: f32[9], param_1.157: s32[256], param_2.101: f32[256], param_3.3099: token[]) -> f32[9] { + %param_3.3099 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.105 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.157 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.101 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -998,47 +998,47 @@ StackFrames ROOT %scatter_offload_custom_fusion.54.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.54.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.7 (param_0.106: f32[9], param_1.158: s32[256], param_2.102: f32[256], param_3.3100: token[]) -> f32[9] { - %param_3.3100 = token[] parameter(3) +%async_computation.7 (param_0.106: f32[9], param_1.158: s32[256], param_2.102: f32[256], param_3.3098: token[]) -> f32[9] { + %param_3.3098 = token[] parameter(3) %param_0.106 = f32[9]{0:T(128)} parameter(0) %param_1.158 = s32[256]{0:T(256)} parameter(1) %param_2.102 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.23.cloned.1 = f32[9]{0:T(128)} call(%param_0.106, %param_1.158, %param_2.102, %param_3.3100), to_apply=%called_computation.7, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.23.cloned.1 = f32[9]{0:T(128)} call(%param_0.106, %param_1.158, %param_2.102, %param_3.3098), to_apply=%called_computation.7, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_104.122 (scatter-add.83: s32[], scatter-add.84: s32[]) -> s32[] { %scatter-add.83 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.84 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1410 = s32[]{:T(128)S(7)} add(%scatter-add.83, %scatter-add.84), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.31.clone.clone (param_0.4565: s32[263], param_1.5349: s32[8], param_2.4518: s32[8]) -> s32[263] { - %param_0.4565 = s32[263]{0:T(512)} parameter(0) - %param_1.5349 = s32[8]{0:T(128)} parameter(1) - %reshape.3919 = s32[8]{0:T(128)} reshape(%param_1.5349), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %transpose.1116 = s32[8]{0:T(128)} transpose(%reshape.3919), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %param_2.4518 = s32[8]{0:T(128)} parameter(2) - %reshape.3920 = s32[8]{0:T(128)} reshape(%param_2.4518), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1117 = s32[8]{0:T(128)} transpose(%reshape.3920), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.239 = s32[263]{0:T(512)} scatter(%param_0.4565, %transpose.1116, %transpose.1117), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_104.122, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %add.1400 = s32[]{:T(128)S(7)} add(%scatter-add.83, %scatter-add.84), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.30 (param_0.4566: s32[263], param_1.5350: s32[8], param_2.4519: s32[8]) -> s32[263] { +%fused_computation.31.clone.clone (param_0.4566: s32[263], param_1.5345: s32[8], param_2.4516: s32[8]) -> s32[263] { %param_0.4566 = s32[263]{0:T(512)} parameter(0) - %param_1.5350 = s32[8]{0:T(128)} parameter(1) - %param_2.4519 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.55 = s32[263]{0:T(512)} fusion(%param_0.4566, %param_1.5350, %param_2.4519), kind=kCustom, calls=%fused_computation.31.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5345 = s32[8]{0:T(128)} parameter(1) + %reshape.3939 = s32[8]{0:T(128)} reshape(%param_1.5345), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %transpose.1116 = s32[8]{0:T(128)} transpose(%reshape.3939), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %param_2.4516 = s32[8]{0:T(128)} parameter(2) + %reshape.3940 = s32[8]{0:T(128)} reshape(%param_2.4516), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %transpose.1117 = s32[8]{0:T(128)} transpose(%reshape.3940), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + ROOT %scatter-add.245 = s32[263]{0:T(512)} scatter(%param_0.4566, %transpose.1116, %transpose.1117), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_104.122, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.30 (param_0.4567: s32[263], param_1.5351: s32[8], param_2.4520: s32[8]) -> s32[263] { +%called_computation.30 (param_0.4567: s32[263], param_1.5346: s32[8], param_2.4517: s32[8]) -> s32[263] { %param_0.4567 = s32[263]{0:T(512)} parameter(0) - %param_1.5351 = s32[8]{0:T(128)} parameter(1) - %param_2.4520 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.56.cloned.1 = s32[263]{0:T(512)} call(%param_0.4567, %param_1.5351, %param_2.4520), to_apply=%called_computation.30, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + %param_1.5346 = s32[8]{0:T(128)} parameter(1) + %param_2.4517 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.55 = s32[263]{0:T(512)} fusion(%param_0.4567, %param_1.5346, %param_2.4517), kind=kCustom, calls=%fused_computation.31.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.8 (param_0.108: s32[263], param_1.160: s32[8], param_2.104: s32[8], param_3.3109: token[]) -> s32[263] { - %param_3.3109 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.30 (param_0.4568: s32[263], param_1.5347: s32[8], param_2.4518: s32[8]) -> s32[263] { + %param_0.4568 = s32[263]{0:T(512)} parameter(0) + %param_1.5347 = s32[8]{0:T(128)} parameter(1) + %param_2.4518 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.56.cloned.1 = s32[263]{0:T(512)} call(%param_0.4568, %param_1.5347, %param_2.4518), to_apply=%called_computation.30, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.8 (param_0.108: s32[263], param_1.160: s32[8], param_2.104: s32[8], param_3.3107: token[]) -> s32[263] { + %param_3.3107 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.108 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.160 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.104 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1046,47 +1046,47 @@ StackFrames ROOT %scatter_offload_custom_fusion.56.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.56.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.8 (param_0.109: s32[263], param_1.161: s32[8], param_2.105: s32[8], param_3.3108: token[]) -> s32[263] { - %param_3.3108 = token[] parameter(3) +%async_computation.8 (param_0.109: s32[263], param_1.161: s32[8], param_2.105: s32[8], param_3.3106: token[]) -> s32[263] { + %param_3.3106 = token[] parameter(3) %param_0.109 = s32[263]{0:T(512)} parameter(0) %param_1.161 = s32[8]{0:T(128)} parameter(1) %param_2.105 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.26.cloned.1 = s32[263]{0:T(512)} call(%param_0.109, %param_1.161, %param_2.105, %param_3.3108), to_apply=%called_computation.8, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.26.cloned.1 = s32[263]{0:T(512)} call(%param_0.109, %param_1.161, %param_2.105, %param_3.3106), to_apply=%called_computation.8, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_14.20 (scatter-add.0: s32[], scatter-add.1: s32[]) -> s32[] { %scatter-add.0 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.1 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1322 = s32[]{:T(128)S(7)} add(%scatter-add.0, %scatter-add.1), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.17.clone.clone.clone (param_0.4570: s32[256], param_1.5352: s32[4096], param_2.4521: s32[4096]) -> s32[256] { - %param_0.4570 = s32[256]{0:T(256)} parameter(0) - %param_1.5352 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3921 = s32[4096]{0:T(1024)} reshape(%param_1.5352), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} - %transpose.1118 = s32[4096]{0:T(1024)} transpose(%reshape.3921), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} - %param_2.4521 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3922 = s32[4096]{0:T(1024)} reshape(%param_2.4521), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - %transpose.1119 = s32[4096]{0:T(1024)} transpose(%reshape.3922), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.240 = s32[256]{0:T(256)} scatter(%param_0.4570, %transpose.1118, %transpose.1119), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_14.20, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %add.1312 = s32[]{:T(128)S(7)} add(%scatter-add.0, %scatter-add.1), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.32 (param_0.4571: s32[256], param_1.5353: s32[4096], param_2.4522: s32[4096]) -> s32[256] { +%fused_computation.17.clone.clone.clone (param_0.4571: s32[256], param_1.5348: s32[4096], param_2.4519: s32[4096]) -> s32[256] { %param_0.4571 = s32[256]{0:T(256)} parameter(0) - %param_1.5353 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4522 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.57 = s32[256]{0:T(256)} fusion(%param_0.4571, %param_1.5353, %param_2.4522), kind=kCustom, calls=%fused_computation.17.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5348 = s32[4096]{0:T(1024)} parameter(1) + %reshape.3941 = s32[4096]{0:T(1024)} reshape(%param_1.5348), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} + %transpose.1118 = s32[4096]{0:T(1024)} transpose(%reshape.3941), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} + %param_2.4519 = s32[4096]{0:T(1024)} parameter(2) + %reshape.3942 = s32[4096]{0:T(1024)} reshape(%param_2.4519), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + %transpose.1119 = s32[4096]{0:T(1024)} transpose(%reshape.3942), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.246 = s32[256]{0:T(256)} scatter(%param_0.4571, %transpose.1118, %transpose.1119), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_14.20, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.32 (param_0.4572: s32[256], param_1.5354: s32[4096], param_2.4523: s32[4096]) -> s32[256] { +%called_computation.32 (param_0.4572: s32[256], param_1.5349: s32[4096], param_2.4520: s32[4096]) -> s32[256] { %param_0.4572 = s32[256]{0:T(256)} parameter(0) - %param_1.5354 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4523 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.58.cloned.1 = s32[256]{0:T(256)} call(%param_0.4572, %param_1.5354, %param_2.4523), to_apply=%called_computation.32, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} + %param_1.5349 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4520 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.57 = s32[256]{0:T(256)} fusion(%param_0.4572, %param_1.5349, %param_2.4520), kind=kCustom, calls=%fused_computation.17.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.9 (param_0.111: s32[256], param_1.163: s32[4096], param_2.107: s32[4096], param_3.3087: token[]) -> s32[256] { - %param_3.3087 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.32 (param_0.4573: s32[256], param_1.5350: s32[4096], param_2.4521: s32[4096]) -> s32[256] { + %param_0.4573 = s32[256]{0:T(256)} parameter(0) + %param_1.5350 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4521 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.58.cloned.1 = s32[256]{0:T(256)} call(%param_0.4573, %param_1.5350, %param_2.4521), to_apply=%called_computation.32, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.9 (param_0.111: s32[256], param_1.163: s32[4096], param_2.107: s32[4096], param_3.3085: token[]) -> s32[256] { + %param_3.3085 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.111 = s32[256]{0:T(256)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.163 = s32[4096]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.107 = s32[4096]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1094,57 +1094,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.58.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.58.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.9 (param_0.112: s32[256], param_1.164: s32[4096], param_2.108: s32[4096], param_3.3086: token[]) -> s32[256] { - %param_3.3086 = token[] parameter(3) +%async_computation.9 (param_0.112: s32[256], param_1.164: s32[4096], param_2.108: s32[4096], param_3.3084: token[]) -> s32[256] { + %param_3.3084 = token[] parameter(3) %param_0.112 = s32[256]{0:T(256)} parameter(0) %param_1.164 = s32[4096]{0:T(1024)} parameter(1) %param_2.108 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.29.cloned.1 = s32[256]{0:T(256)} call(%param_0.112, %param_1.164, %param_2.108, %param_3.3086), to_apply=%called_computation.9, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.29.cloned.1 = s32[256]{0:T(256)} call(%param_0.112, %param_1.164, %param_2.108, %param_3.3084), to_apply=%called_computation.9, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.33 (param_0.4573: s32[263]) -> s32[263] { - %param_0.4573 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2093 = s32[263]{0:T(512)} copy(%param_0.4573), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.33 (param_0.4574: s32[263]) -> s32[263] { + %param_0.4574 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2093 = s32[263]{0:T(512)} copy(%param_0.4574), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.33 (param_0.4574: s32[263]) -> s32[263] { - %param_0.4574 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2094.cloned.1 = s32[263]{0:T(512)} call(%param_0.4574), to_apply=%called_computation.33 +%async_computation.33 (param_0.4575: s32[263]) -> s32[263] { + %param_0.4575 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2094.cloned.1 = s32[263]{0:T(512)} call(%param_0.4575), to_apply=%called_computation.33 }, execution_thread="sparsecore" %region_20.26.clone.1 (scatter-add.141: s32[], scatter-add.142: s32[]) -> s32[] { %scatter-add.141 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.142 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2471 = s32[]{:T(128)S(7)} add(%scatter-add.141, %scatter-add.142), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.18.clone.clone.clone (param_0.4575: s32[263], param_1.5355: s32[256], param_2.4524: s32[256]) -> s32[263] { - %param_0.4575 = s32[263]{0:T(512)} parameter(0) - %param_1.5355 = s32[256]{0:T(256)} parameter(1) - %reshape.3923 = s32[256]{0:T(256)} reshape(%param_1.5355), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1120 = s32[256]{0:T(256)} transpose(%reshape.3923), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4524 = s32[256]{0:T(256)} parameter(2) - %reshape.3924 = s32[256]{0:T(256)} reshape(%param_2.4524) - %transpose.1121 = s32[256]{0:T(256)} transpose(%reshape.3924), dimensions={0} - ROOT %scatter-add.241 = s32[263]{0:T(512)} scatter(%param_0.4575, %transpose.1120, %transpose.1121), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_20.26.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %add.2463 = s32[]{:T(128)S(7)} add(%scatter-add.141, %scatter-add.142), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.34 (param_0.4576: s32[263], param_1.5356: s32[256], param_2.4525: s32[256]) -> s32[263] { +%fused_computation.18.clone.clone.clone (param_0.4576: s32[263], param_1.5351: s32[256], param_2.4522: s32[256]) -> s32[263] { %param_0.4576 = s32[263]{0:T(512)} parameter(0) - %param_1.5356 = s32[256]{0:T(256)} parameter(1) - %param_2.4525 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.59 = s32[263]{0:T(512)} fusion(%param_0.4576, %param_1.5356, %param_2.4525), kind=kCustom, calls=%fused_computation.18.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5351 = s32[256]{0:T(256)} parameter(1) + %reshape.3943 = s32[256]{0:T(256)} reshape(%param_1.5351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1120 = s32[256]{0:T(256)} transpose(%reshape.3943), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4522 = s32[256]{0:T(256)} parameter(2) + %reshape.3944 = s32[256]{0:T(256)} reshape(%param_2.4522) + %transpose.1121 = s32[256]{0:T(256)} transpose(%reshape.3944), dimensions={0} + ROOT %scatter-add.247 = s32[263]{0:T(512)} scatter(%param_0.4576, %transpose.1120, %transpose.1121), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_20.26.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.34 (param_0.4577: s32[263], param_1.5357: s32[256], param_2.4526: s32[256]) -> s32[263] { +%called_computation.34 (param_0.4577: s32[263], param_1.5352: s32[256], param_2.4523: s32[256]) -> s32[263] { %param_0.4577 = s32[263]{0:T(512)} parameter(0) - %param_1.5357 = s32[256]{0:T(256)} parameter(1) - %param_2.4526 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.60.cloned.1 = s32[263]{0:T(512)} call(%param_0.4577, %param_1.5357, %param_2.4526), to_apply=%called_computation.34, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5352 = s32[256]{0:T(256)} parameter(1) + %param_2.4523 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.59 = s32[263]{0:T(512)} fusion(%param_0.4577, %param_1.5352, %param_2.4523), kind=kCustom, calls=%fused_computation.18.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.10 (param_0.114: s32[263], param_1.166: s32[256], param_2.110: s32[256], param_3.3089: token[]) -> s32[263] { - %param_3.3089 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.34 (param_0.4578: s32[263], param_1.5353: s32[256], param_2.4524: s32[256]) -> s32[263] { + %param_0.4578 = s32[263]{0:T(512)} parameter(0) + %param_1.5353 = s32[256]{0:T(256)} parameter(1) + %param_2.4524 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.60.cloned.1 = s32[263]{0:T(512)} call(%param_0.4578, %param_1.5353, %param_2.4524), to_apply=%called_computation.34, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.10 (param_0.114: s32[263], param_1.166: s32[256], param_2.110: s32[256], param_3.3087: token[]) -> s32[263] { + %param_3.3087 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.114 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.166 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.110 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1154,57 +1154,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.60.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.60.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.10 (param_0.115: s32[263], param_1.167: s32[256], param_2.111: s32[256], param_3.3088: token[]) -> s32[263] { - %param_3.3088 = token[] parameter(3) +%async_computation.10 (param_0.115: s32[263], param_1.167: s32[256], param_2.111: s32[256], param_3.3086: token[]) -> s32[263] { + %param_3.3086 = token[] parameter(3) %param_0.115 = s32[263]{0:T(512)} parameter(0) %param_1.167 = s32[256]{0:T(256)} parameter(1) %param_2.111 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.32.cloned.1 = s32[263]{0:T(512)} call(%param_0.115, %param_1.167, %param_2.111, %param_3.3088), to_apply=%called_computation.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.32.cloned.1 = s32[263]{0:T(512)} call(%param_0.115, %param_1.167, %param_2.111, %param_3.3086), to_apply=%called_computation.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.35 (param_0.4578: f32[9]) -> f32[9] { - %param_0.4578 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2095 = f32[9]{0:T(128)} copy(%param_0.4578), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.35 (param_0.4579: f32[9]) -> f32[9] { + %param_0.4579 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2095 = f32[9]{0:T(128)} copy(%param_0.4579), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.35 (param_0.4579: f32[9]) -> f32[9] { - %param_0.4579 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2096.cloned.1 = f32[9]{0:T(128)} call(%param_0.4579), to_apply=%called_computation.35 +%async_computation.35 (param_0.4580: f32[9]) -> f32[9] { + %param_0.4580 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2096.cloned.1 = f32[9]{0:T(128)} call(%param_0.4580), to_apply=%called_computation.35 }, execution_thread="sparsecore" %region_26.33.clone.1 (scatter-add.145: f32[], scatter-add.146: f32[]) -> f32[] { %scatter-add.145 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.146 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2473 = f32[]{:T(128)S(7)} add(%scatter-add.145, %scatter-add.146), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.19.clone.clone.clone (param_0.4580: f32[9], param_1.5358: s32[256], param_2.4527: f32[256]) -> f32[9] { - %param_0.4580 = f32[9]{0:T(128)} parameter(0) - %param_1.5358 = s32[256]{0:T(256)} parameter(1) - %reshape.3925 = s32[256]{0:T(256)} reshape(%param_1.5358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1122 = s32[256]{0:T(256)} transpose(%reshape.3925), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4527 = f32[256]{0:T(256)} parameter(2) - %reshape.3926 = f32[256]{0:T(256)} reshape(%param_2.4527) - %transpose.1123 = f32[256]{0:T(256)} transpose(%reshape.3926), dimensions={0} - ROOT %scatter-add.242 = f32[9]{0:T(128)} scatter(%param_0.4580, %transpose.1122, %transpose.1123), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_26.33.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %add.2465 = f32[]{:T(128)S(7)} add(%scatter-add.145, %scatter-add.146), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.36 (param_0.4581: f32[9], param_1.5359: s32[256], param_2.4528: f32[256]) -> f32[9] { +%fused_computation.19.clone.clone.clone (param_0.4581: f32[9], param_1.5354: s32[256], param_2.4525: f32[256]) -> f32[9] { %param_0.4581 = f32[9]{0:T(128)} parameter(0) - %param_1.5359 = s32[256]{0:T(256)} parameter(1) - %param_2.4528 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.61 = f32[9]{0:T(128)} fusion(%param_0.4581, %param_1.5359, %param_2.4528), kind=kCustom, calls=%fused_computation.19.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5354 = s32[256]{0:T(256)} parameter(1) + %reshape.3945 = s32[256]{0:T(256)} reshape(%param_1.5354), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1122 = s32[256]{0:T(256)} transpose(%reshape.3945), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4525 = f32[256]{0:T(256)} parameter(2) + %reshape.3946 = f32[256]{0:T(256)} reshape(%param_2.4525) + %transpose.1123 = f32[256]{0:T(256)} transpose(%reshape.3946), dimensions={0} + ROOT %scatter-add.248 = f32[9]{0:T(128)} scatter(%param_0.4581, %transpose.1122, %transpose.1123), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_26.33.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.36 (param_0.4582: f32[9], param_1.5360: s32[256], param_2.4529: f32[256]) -> f32[9] { +%called_computation.36 (param_0.4582: f32[9], param_1.5355: s32[256], param_2.4526: f32[256]) -> f32[9] { %param_0.4582 = f32[9]{0:T(128)} parameter(0) - %param_1.5360 = s32[256]{0:T(256)} parameter(1) - %param_2.4529 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.62.cloned.1 = f32[9]{0:T(128)} call(%param_0.4582, %param_1.5360, %param_2.4529), to_apply=%called_computation.36, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5355 = s32[256]{0:T(256)} parameter(1) + %param_2.4526 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.61 = f32[9]{0:T(128)} fusion(%param_0.4582, %param_1.5355, %param_2.4526), kind=kCustom, calls=%fused_computation.19.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.11 (param_0.117: f32[9], param_1.169: s32[256], param_2.113: f32[256], param_3.3095: token[]) -> f32[9] { - %param_3.3095 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.36 (param_0.4583: f32[9], param_1.5356: s32[256], param_2.4527: f32[256]) -> f32[9] { + %param_0.4583 = f32[9]{0:T(128)} parameter(0) + %param_1.5356 = s32[256]{0:T(256)} parameter(1) + %param_2.4527 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.62.cloned.1 = f32[9]{0:T(128)} call(%param_0.4583, %param_1.5356, %param_2.4527), to_apply=%called_computation.36, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.11 (param_0.117: f32[9], param_1.169: s32[256], param_2.113: f32[256], param_3.3093: token[]) -> f32[9] { + %param_3.3093 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.117 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.169 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.113 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1214,57 +1214,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.62.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.62.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.11 (param_0.118: f32[9], param_1.170: s32[256], param_2.114: f32[256], param_3.3094: token[]) -> f32[9] { - %param_3.3094 = token[] parameter(3) +%async_computation.11 (param_0.118: f32[9], param_1.170: s32[256], param_2.114: f32[256], param_3.3092: token[]) -> f32[9] { + %param_3.3092 = token[] parameter(3) %param_0.118 = f32[9]{0:T(128)} parameter(0) %param_1.170 = s32[256]{0:T(256)} parameter(1) %param_2.114 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.35.cloned.1 = f32[9]{0:T(128)} call(%param_0.118, %param_1.170, %param_2.114, %param_3.3094), to_apply=%called_computation.11, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.35.cloned.1 = f32[9]{0:T(128)} call(%param_0.118, %param_1.170, %param_2.114, %param_3.3092), to_apply=%called_computation.11, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.37 (param_0.4583: s32[263]) -> s32[263] { - %param_0.4583 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2097 = s32[263]{0:T(512)} copy(%param_0.4583), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.37 (param_0.4584: s32[263]) -> s32[263] { + %param_0.4584 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2097 = s32[263]{0:T(512)} copy(%param_0.4584), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.37 (param_0.4584: s32[263]) -> s32[263] { - %param_0.4584 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2098.cloned.1 = s32[263]{0:T(512)} call(%param_0.4584), to_apply=%called_computation.37 +%async_computation.37 (param_0.4585: s32[263]) -> s32[263] { + %param_0.4585 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2098.cloned.1 = s32[263]{0:T(512)} call(%param_0.4585), to_apply=%called_computation.37 }, execution_thread="sparsecore" %region_28.35.clone.1 (scatter-add.149: s32[], scatter-add.150: s32[]) -> s32[] { %scatter-add.149 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.150 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2475 = s32[]{:T(128)S(7)} add(%scatter-add.149, %scatter-add.150), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.2467 = s32[]{:T(128)S(7)} add(%scatter-add.149, %scatter-add.150), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.20.clone.clone.clone (param_0.4585: s32[263], param_1.5361: s32[8], param_2.4530: s32[8]) -> s32[263] { - %param_0.4585 = s32[263]{0:T(512)} parameter(0) - %param_1.5361 = s32[8]{0:T(128)} parameter(1) - %reshape.3927 = s32[8]{0:T(128)} reshape(%param_1.5361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1124 = s32[8]{0:T(128)} transpose(%reshape.3927), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4530 = s32[8]{0:T(128)} parameter(2) - %reshape.3928 = s32[8]{0:T(128)} reshape(%param_2.4530) - %transpose.1125 = s32[8]{0:T(128)} transpose(%reshape.3928), dimensions={0} - ROOT %scatter-add.243 = s32[263]{0:T(512)} scatter(%param_0.4585, %transpose.1124, %transpose.1125), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_28.35.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.38 (param_0.4586: s32[263], param_1.5362: s32[8], param_2.4531: s32[8]) -> s32[263] { +%fused_computation.20.clone.clone.clone (param_0.4586: s32[263], param_1.5357: s32[8], param_2.4528: s32[8]) -> s32[263] { %param_0.4586 = s32[263]{0:T(512)} parameter(0) - %param_1.5362 = s32[8]{0:T(128)} parameter(1) - %param_2.4531 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.63 = s32[263]{0:T(512)} fusion(%param_0.4586, %param_1.5362, %param_2.4531), kind=kCustom, calls=%fused_computation.20.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5357 = s32[8]{0:T(128)} parameter(1) + %reshape.3947 = s32[8]{0:T(128)} reshape(%param_1.5357), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1124 = s32[8]{0:T(128)} transpose(%reshape.3947), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4528 = s32[8]{0:T(128)} parameter(2) + %reshape.3948 = s32[8]{0:T(128)} reshape(%param_2.4528) + %transpose.1125 = s32[8]{0:T(128)} transpose(%reshape.3948), dimensions={0} + ROOT %scatter-add.249 = s32[263]{0:T(512)} scatter(%param_0.4586, %transpose.1124, %transpose.1125), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_28.35.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.38 (param_0.4587: s32[263], param_1.5363: s32[8], param_2.4532: s32[8]) -> s32[263] { +%called_computation.38 (param_0.4587: s32[263], param_1.5358: s32[8], param_2.4529: s32[8]) -> s32[263] { %param_0.4587 = s32[263]{0:T(512)} parameter(0) - %param_1.5363 = s32[8]{0:T(128)} parameter(1) - %param_2.4532 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.64.cloned.1 = s32[263]{0:T(512)} call(%param_0.4587, %param_1.5363, %param_2.4532), to_apply=%called_computation.38, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5358 = s32[8]{0:T(128)} parameter(1) + %param_2.4529 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.63 = s32[263]{0:T(512)} fusion(%param_0.4587, %param_1.5358, %param_2.4529), kind=kCustom, calls=%fused_computation.20.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.12 (param_0.120: s32[263], param_1.172: s32[8], param_2.116: s32[8], param_3.3103: token[]) -> s32[263] { - %param_3.3103 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.38 (param_0.4588: s32[263], param_1.5359: s32[8], param_2.4530: s32[8]) -> s32[263] { + %param_0.4588 = s32[263]{0:T(512)} parameter(0) + %param_1.5359 = s32[8]{0:T(128)} parameter(1) + %param_2.4530 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.64.cloned.1 = s32[263]{0:T(512)} call(%param_0.4588, %param_1.5359, %param_2.4530), to_apply=%called_computation.38, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.12 (param_0.120: s32[263], param_1.172: s32[8], param_2.116: s32[8], param_3.3101: token[]) -> s32[263] { + %param_3.3101 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.120 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.172 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.116 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1274,12 +1274,12 @@ StackFrames ROOT %scatter_offload_custom_fusion.64.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.64.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.12 (param_0.121: s32[263], param_1.173: s32[8], param_2.117: s32[8], param_3.3102: token[]) -> s32[263] { - %param_3.3102 = token[] parameter(3) +%async_computation.12 (param_0.121: s32[263], param_1.173: s32[8], param_2.117: s32[8], param_3.3100: token[]) -> s32[263] { + %param_3.3100 = token[] parameter(3) %param_0.121 = s32[263]{0:T(512)} parameter(0) %param_1.173 = s32[8]{0:T(128)} parameter(1) %param_2.117 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.38.cloned.1 = s32[263]{0:T(512)} call(%param_0.121, %param_1.173, %param_2.117, %param_3.3102), to_apply=%called_computation.12, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.38.cloned.1 = s32[263]{0:T(512)} call(%param_0.121, %param_1.173, %param_2.117, %param_3.3100), to_apply=%called_computation.12, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_154.179 (reduce_sum.431: f32[], reduce_sum.254: f32[]) -> f32[] { @@ -1288,18 +1288,18 @@ StackFrames ROOT %reduce_sum.258 = f32[]{:T(128)} add(%reduce_sum.431, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.467 (param_0.4170: f32[3,1536,128,192]) -> f32[] { - %param_0.4170 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(0) - %bitcast.672 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4170), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.564 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.672, %bitcast.672), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5086 = f32[]{:T(128)} constant(0) - ROOT %reduce.669 = f32[]{:T(128)} reduce(%square.564, %constant.5086), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.466 (param_0.4171: f32[3,1536,128,192]) -> f32[] { + %param_0.4171 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(0) + %bitcast.670 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4171), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3798 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.670, %bitcast.670), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5104 = f32[]{:T(128)} constant(0) + ROOT %reduce.669 = f32[]{:T(128)} reduce(%mul.3798, %constant.5104), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.468 (param_0.1421: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { - %param_0.1421 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) - %copy.1550 = bf16[1536,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1421), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\']"} - ROOT %bitcast.673 = bf16[3,1536,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1550), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} +%fused_computation.467 (param_0.1419: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { + %param_0.1419 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %copy.1550 = bf16[1536,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1419), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\']"} + ROOT %bitcast.671 = bf16[3,1536,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1550), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} } %region_221.246 (reduce_sum.893: f32[], reduce_sum.603: f32[]) -> f32[] { @@ -1314,55 +1314,55 @@ StackFrames ROOT %reduce_sum.450 = f32[]{:T(128)} add(%reduce_sum.655, %reduce_sum.449), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.469 (param_0.4140: f32[1536,3,128,192], param_1.5025: f32[], param_2.4298: f32[], param_3.2951: f32[], param_4.2203: f32[1536,3,128,192], param_5.2006: f32[], param_6.1443: f32[3,1536,128,192], param_7.1124: pred[], param_8.889: f32[1536,3,128,192]) -> (f32[], f32[1536,3,128,192], f32[1536,3,128,192], f32[1536,3,128,192], f32[]) { - %param_0.4140 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) - %param_3.2951 = f32[]{:T(128)S(6)} parameter(3) - %mul.4715.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.468 (param_0.4141: f32[1536,3,128,192], param_1.5021: f32[], param_2.4296: f32[], param_3.2949: f32[], param_4.2203: f32[1536,3,128,192], param_5.2006: f32[], param_6.1443: f32[3,1536,128,192], param_7.1124: pred[], param_8.889: f32[1536,3,128,192]) -> (f32[], f32[1536,3,128,192], f32[1536,3,128,192], f32[1536,3,128,192], f32[]) { + %param_0.4141 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %param_3.2949 = f32[]{:T(128)S(6)} parameter(3) + %mul.5043.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2949), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1124 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2121.clone.1 = pred[1536,3,128,192]{2,3,1,0:T(8,128)(4,1)} broadcast(%param_7.1124), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2165.clone.1 = pred[1536,3,128,192]{2,3,1,0:T(8,128)(4,1)} broadcast(%param_7.1124), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1443 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(6) - %bitcast.1374.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_6.1443), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1372.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_6.1443), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} %param_5.2006 = f32[]{:T(128)} parameter(5) - %div.2565.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_5.2006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2564.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.1374.clone.1, %div.2565.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2120.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2121.clone.1, %bitcast.1374.clone.1, %div.2564.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4845.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4252.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4845.clone.1), dimensions={}, metadata={op_name="broadcast.334"} - %mul.4721.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2575.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_5.2006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2574.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.1372.clone.1, %div.2575.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2164.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2165.clone.1, %bitcast.1372.clone.1, %div.2574.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4863.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4279.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4863.clone.1), dimensions={}, metadata={op_name="broadcast.334"} + %mul.5049.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2164.clone.1, %broadcast.4279.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.889 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(8) - %constant.4849.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4722.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4849.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4720.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.4722.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3429.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4721.clone.1, %mul.4720.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4298 = f32[]{:T(128)S(6)} parameter(2) - %div.2561.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4298), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.399.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %select_n.2120.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4848.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4719.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4848.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4717.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.4719.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4867.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.5050.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4867.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5048.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.5050.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3443.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.5049.clone.1, %mul.5048.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4296 = f32[]{:T(128)S(6)} parameter(2) + %div.2571.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4296), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.399.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2164.clone.1, %select_n.2164.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4866.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.5047.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4866.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5045.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.5047.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2203 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(4) - %constant.4847.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4718.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4847.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4716.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.4718.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3428.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4717.clone.1, %mul.4716.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5025 = f32[]{:T(128)S(6)} parameter(1) - %div.2560.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5025), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2559.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3428.clone.1, %div.2560.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.157.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} sqrt(%div.2559.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4846.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3427.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4846.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3426.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.157.clone.1, %add.3427.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1293.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%div.2561.clone.1, %add.3426.clone.1), metadata={op_name="multiply.290"} - %div.2558.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3429.clone.1, %multiply.1293.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4714.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4140, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3425.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2558.clone.1, %mul.4714.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4713.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.4715.clone.1, %add.3425.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3424.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4140, %mul.4713.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.565 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3424.clone.1, %add.3424.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5056 = f32[]{:T(128)} constant(0) - %reduce.670 = f32[]{:T(128)} reduce(%square.565, %constant.5056), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.671.clone.1 = f32[]{:T(128)} reduce(%integer_pow.399.clone.1, %constant.5056), dimensions={0,1,2,3}, to_apply=%region_187.212, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.656 = (f32[]{:T(128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.670, %add.3424.clone.1, %add.3428.clone.1, %add.3429.clone.1, %reduce.671.clone.1) + %constant.4865.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.5046.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4865.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5044.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.5046.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3442.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.5045.clone.1, %mul.5044.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5021 = f32[]{:T(128)S(6)} parameter(1) + %div.2570.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5021), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2569.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3442.clone.1, %div.2570.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.157.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} sqrt(%div.2569.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4864.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3441.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4864.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3440.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.157.clone.1, %add.3441.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1293.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%div.2571.clone.1, %add.3440.clone.1), metadata={op_name="multiply.290"} + %div.2568.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3443.clone.1, %multiply.1293.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.5042.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4141, %broadcast.4279.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3439.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2568.clone.1, %mul.5042.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.5041.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.5043.clone.1, %add.3439.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3438.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4141, %mul.5041.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.330 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3438.clone.1, %add.3438.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5074 = f32[]{:T(128)} constant(0) + %reduce.670 = f32[]{:T(128)} reduce(%square.330, %constant.5074), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.671.clone.1 = f32[]{:T(128)} reduce(%integer_pow.399.clone.1, %constant.5074), dimensions={0,1,2,3}, to_apply=%region_187.212, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.660 = (f32[]{:T(128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.670, %add.3438.clone.1, %add.3442.clone.1, %add.3443.clone.1, %reduce.671.clone.1) } %region_160.185 (reduce_sum.473: f32[], reduce_sum.293: f32[]) -> f32[] { @@ -1377,19 +1377,19 @@ StackFrames ROOT %reduce_sum.461 = f32[]{:T(128)} add(%reduce_sum.459, %reduce_sum.460), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.495 (param_0.4166: bf16[256,512,512], param_1.5047: bf16[256,512,512]) -> (f32[], f32[]) { - %param_0.4166 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) - %broadcast_in_dim.1245 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4166), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.695 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.570 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.695, %bitcast.695), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5082 = f32[]{:T(128)} constant(0) - %reduce.672 = f32[]{:T(128)} reduce(%square.570, %constant.5082), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.5047 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(1) - %broadcast_in_dim.1253.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_1.5047), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.703.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1253.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.576.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.703.clone.1, %bitcast.703.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.674.clone.1 = f32[]{:T(128)} reduce(%square.576.clone.1, %constant.5082), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.763 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.672, %reduce.674.clone.1) +%fused_computation.494 (param_0.4167: bf16[256,512,512], param_1.5043: bf16[256,512,512]) -> (f32[], f32[]) { + %param_0.4167 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %broadcast_in_dim.1358 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4167), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.693 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3827 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.693, %bitcast.693), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5100 = f32[]{:T(128)} constant(0) + %reduce.672 = f32[]{:T(128)} reduce(%mul.3827, %constant.5100), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.5043 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(1) + %broadcast_in_dim.1366.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_1.5043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.701.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1366.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3833.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.701.clone.1, %bitcast.701.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.674.clone.1 = f32[]{:T(128)} reduce(%mul.3833.clone.1, %constant.5100), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.767 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.672, %reduce.674.clone.1) } %region_159.184 (reduce_sum.466: f32[], reduce_sum.279: f32[]) -> f32[] { @@ -1398,13 +1398,13 @@ StackFrames ROOT %reduce_sum.286 = f32[]{:T(128)} add(%reduce_sum.466, %reduce_sum.279), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.497 (param_0.4165: bf16[256,512,512]) -> f32[] { - %param_0.4165 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) - %broadcast_in_dim.1249 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4165), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.699 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.573 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.699, %bitcast.699), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5081 = f32[]{:T(128)} constant(0) - ROOT %reduce.673 = f32[]{:T(128)} reduce(%square.573, %constant.5081), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.496 (param_0.4166: bf16[256,512,512]) -> f32[] { + %param_0.4166 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %broadcast_in_dim.1362 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4166), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.697 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3830 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.697, %bitcast.697), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5099 = f32[]{:T(128)} constant(0) + ROOT %reduce.673 = f32[]{:T(128)} reduce(%mul.3830, %constant.5099), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_227.252 (reduce_sum.935: f32[], reduce_sum.631: f32[]) -> f32[] { @@ -1419,61 +1419,61 @@ StackFrames ROOT %reduce_sum.472 = f32[]{:T(128)} add(%reduce_sum.697, %reduce_sum.471), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.515 (param_0.4134: f32[], param_1.5019: f32[256,1,512,512], param_2.4292: f32[], param_3.2945: f32[256,1,512,512], param_4.2197: f32[], param_5.2000: bf16[256,512,512], param_6.1437: pred[], param_7.1118: f32[], param_8.883: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.514 (param_0.4135: f32[], param_1.5015: f32[256,1,512,512], param_2.4290: f32[], param_3.2943: f32[256,1,512,512], param_4.2197: f32[], param_5.2000: bf16[256,512,512], param_6.1437: pred[], param_7.1118: f32[], param_8.883: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.883 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1359.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.883), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %bitcast.1357.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.883), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} %param_7.1118 = f32[]{:T(128)S(6)} parameter(7) - %mul.4664.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4992.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1437 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2103.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1437), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2147.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1437), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2000 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1459.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2000), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1361.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1459.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %broadcast_in_dim.1572.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2000), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1359.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1572.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2197 = f32[]{:T(128)} parameter(4) - %div.2523.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2197), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2522.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1361.clone.1, %div.2523.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2102.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2103.clone.1, %bitcast.1361.clone.1, %div.2522.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4815.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4232.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4815.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} - %mul.4666.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %broadcast.4232.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2945 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1360.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2945), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %constant.4814.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4231.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4814.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4665.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4231.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4666.clone.1, %mul.4665.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4292 = f32[]{:T(128)S(6)} parameter(2) - %div.2521.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4292), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %select_n.2102.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4813.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4234.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4813.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} - %mul.4668.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4234.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5019 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1362.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5019), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %constant.4812.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4233.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4812.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4667.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4233.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4668.clone.1, %mul.4667.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4134 = f32[]{:T(128)S(6)} parameter(0) - %div.2520.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4134), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2519.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3395.clone.1, %div.2520.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.151.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2519.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4816.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4230.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4816.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.151.clone.1, %broadcast.4230.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1287.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2521.clone.1, %add.3393.clone.1), metadata={op_name="multiply.296"} - %div.2518.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3394.clone.1, %multiply.1287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4663.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1359.clone.1, %broadcast.4232.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3392.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2518.clone.1, %mul.4663.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4662.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4664.clone.1, %add.3392.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3391.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1359.clone.1, %mul.4662.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.577 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3391.clone.1, %add.3391.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5050 = f32[]{:T(128)} constant(0) - %reduce.675 = f32[]{:T(128)} reduce(%square.577, %constant.5050), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.849.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3395.clone.1) - %bitcast.822.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3394.clone.1) - %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5050), dimensions={0,1,2,3}, to_apply=%region_193.218, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.666 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.675, %add.3391.clone.1, %bitcast.849.clone.1, %bitcast.822.clone.1, %reduce.684.clone.1) + %div.2533.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2197), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2532.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1359.clone.1, %div.2533.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2146.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2147.clone.1, %bitcast.1359.clone.1, %div.2532.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4833.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4259.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4833.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} + %mul.4994.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2146.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2943 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1358.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2943), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %constant.4832.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4258.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4832.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.4993.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1358.clone.1, %broadcast.4258.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3408.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4994.clone.1, %mul.4993.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4290 = f32[]{:T(128)S(6)} parameter(2) + %div.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4290), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2146.clone.1, %select_n.2146.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4831.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4261.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4831.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} + %mul.4996.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4261.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5015 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1360.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5015), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %constant.4830.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4260.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4830.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.4995.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4260.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3409.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4996.clone.1, %mul.4995.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4135 = f32[]{:T(128)S(6)} parameter(0) + %div.2530.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4135), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2529.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3409.clone.1, %div.2530.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.151.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2529.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4834.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4257.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4834.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %add.3407.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.151.clone.1, %broadcast.4257.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1287.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2531.clone.1, %add.3407.clone.1), metadata={op_name="multiply.296"} + %div.2528.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3408.clone.1, %multiply.1287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4991.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1357.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3406.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2528.clone.1, %mul.4991.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4990.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4992.clone.1, %add.3406.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3405.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1357.clone.1, %mul.4990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.331 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3405.clone.1, %add.3405.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5068 = f32[]{:T(128)} constant(0) + %reduce.675 = f32[]{:T(128)} reduce(%square.331, %constant.5068), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.847.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3409.clone.1) + %bitcast.820.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3408.clone.1) + %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5068), dimensions={0,1,2,3}, to_apply=%region_193.218, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.670 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.675, %add.3405.clone.1, %bitcast.847.clone.1, %bitcast.820.clone.1, %reduce.684.clone.1) } %region_226.251 (reduce_sum.928: f32[], reduce_sum.625: f32[]) -> f32[] { @@ -1488,61 +1488,61 @@ StackFrames ROOT %reduce_sum.470 = f32[]{:T(128)} add(%reduce_sum.690, %reduce_sum.465), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.516 (param_0.4135: f32[], param_1.5020: f32[256,1,512,512], param_2.4293: f32[], param_3.2946: f32[256,1,512,512], param_4.2198: f32[], param_5.2001: bf16[256,512,512], param_6.1438: pred[], param_7.1119: f32[], param_8.884: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.515 (param_0.4136: f32[], param_1.5016: f32[256,1,512,512], param_2.4291: f32[], param_3.2944: f32[256,1,512,512], param_4.2198: f32[], param_5.2001: bf16[256,512,512], param_6.1438: pred[], param_7.1119: f32[], param_8.884: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.884 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1363.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.884), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %bitcast.1361.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.884), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} %param_7.1119 = f32[]{:T(128)S(6)} parameter(7) - %mul.4671.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4999.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1438 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2105.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1438), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2149.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1438), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2001 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1460.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2001), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1365.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1460.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %broadcast_in_dim.1573.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2001), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1363.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1573.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2198 = f32[]{:T(128)} parameter(4) - %div.2529.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2198), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2528.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1365.clone.1, %div.2529.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2104.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2105.clone.1, %bitcast.1365.clone.1, %div.2528.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4820.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4237.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4820.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} - %mul.4673.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4237.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2946 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1364.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2946), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %constant.4819.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4236.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4819.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4672.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4236.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3399.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4673.clone.1, %mul.4672.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4293 = f32[]{:T(128)S(6)} parameter(2) - %div.2527.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4293), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %select_n.2104.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4818.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4239.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4818.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} - %mul.4675.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4239.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5020 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1366.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5020), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %constant.4817.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4238.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4817.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4674.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4238.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3400.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4675.clone.1, %mul.4674.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4135 = f32[]{:T(128)S(6)} parameter(0) - %div.2526.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4135), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2525.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3400.clone.1, %div.2526.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.152.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2525.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4821.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4235.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4821.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3398.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.152.clone.1, %broadcast.4235.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1288.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2527.clone.1, %add.3398.clone.1), metadata={op_name="multiply.295"} - %div.2524.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3399.clone.1, %multiply.1288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4670.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1363.clone.1, %broadcast.4237.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3397.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2524.clone.1, %mul.4670.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4669.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4671.clone.1, %add.3397.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3396.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1363.clone.1, %mul.4669.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.578 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3396.clone.1, %add.3396.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5051 = f32[]{:T(128)} constant(0) - %reduce.676 = f32[]{:T(128)} reduce(%square.578, %constant.5051), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.840.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3400.clone.1) - %bitcast.813.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3399.clone.1) - %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5051), dimensions={0,1,2,3}, to_apply=%region_192.217, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.665 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.676, %add.3396.clone.1, %bitcast.840.clone.1, %bitcast.813.clone.1, %reduce.685.clone.1) + %div.2539.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2198), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2538.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1363.clone.1, %div.2539.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2148.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2149.clone.1, %bitcast.1363.clone.1, %div.2538.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4838.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4264.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4838.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} + %mul.5001.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2148.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2944 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1362.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2944), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %constant.4837.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4263.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4837.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.5000.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4263.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3413.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.5001.clone.1, %mul.5000.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4291 = f32[]{:T(128)S(6)} parameter(2) + %div.2537.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4291), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2148.clone.1, %select_n.2148.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4836.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4266.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4836.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} + %mul.5003.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4266.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5016 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1364.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5016), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %constant.4835.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4265.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4835.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.5002.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4265.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3414.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.5003.clone.1, %mul.5002.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4136 = f32[]{:T(128)S(6)} parameter(0) + %div.2536.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4136), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2535.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3414.clone.1, %div.2536.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.152.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2535.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4839.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4262.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4839.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %add.3412.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.152.clone.1, %broadcast.4262.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1288.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2537.clone.1, %add.3412.clone.1), metadata={op_name="multiply.295"} + %div.2534.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3413.clone.1, %multiply.1288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4998.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1361.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3411.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2534.clone.1, %mul.4998.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4997.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4999.clone.1, %add.3411.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3410.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1361.clone.1, %mul.4997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.332 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3410.clone.1, %add.3410.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5069 = f32[]{:T(128)} constant(0) + %reduce.676 = f32[]{:T(128)} reduce(%square.332, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.838.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3414.clone.1) + %bitcast.811.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3413.clone.1) + %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_192.217, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.669 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.676, %add.3410.clone.1, %bitcast.838.clone.1, %bitcast.811.clone.1, %reduce.685.clone.1) } %region_225.250 (reduce_sum.921: f32[], reduce_sum.619: f32[]) -> f32[] { @@ -1557,61 +1557,61 @@ StackFrames ROOT %reduce_sum.464 = f32[]{:T(128)} add(%reduce_sum.683, %reduce_sum.463), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.517 (param_0.4136: f32[], param_1.5021: f32[256,1,512,512], param_2.4294: f32[], param_3.2947: f32[256,1,512,512], param_4.2199: f32[], param_5.2002: bf16[256,512,512], param_6.1439: pred[], param_7.1120: f32[], param_8.885: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.516 (param_0.4137: f32[], param_1.5017: f32[256,1,512,512], param_2.4292: f32[], param_3.2945: f32[256,1,512,512], param_4.2199: f32[], param_5.2002: bf16[256,512,512], param_6.1439: pred[], param_7.1120: f32[], param_8.885: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.885 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1367.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.885), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %bitcast.1365.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.885), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} %param_7.1120 = f32[]{:T(128)S(6)} parameter(7) - %mul.4678.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5006.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1439 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2107.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1439), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2151.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1439), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2002 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1461.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2002), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1369.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1461.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %broadcast_in_dim.1574.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2002), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1367.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1574.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2199 = f32[]{:T(128)} parameter(4) - %div.2535.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2199), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2534.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1369.clone.1, %div.2535.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2106.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2107.clone.1, %bitcast.1369.clone.1, %div.2534.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4825.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4242.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4825.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} - %mul.4680.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %broadcast.4242.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2947 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1368.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2947), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4824.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4241.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4824.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4679.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4241.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3404.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4680.clone.1, %mul.4679.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4294 = f32[]{:T(128)S(6)} parameter(2) - %div.2533.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4294), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %select_n.2106.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4823.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4244.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4823.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} - %mul.4682.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4244.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5021 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1370.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5021), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4822.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4243.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4822.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4681.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1370.clone.1, %broadcast.4243.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3405.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4682.clone.1, %mul.4681.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4136 = f32[]{:T(128)S(6)} parameter(0) - %div.2532.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4136), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3405.clone.1, %div.2532.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.153.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2531.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4826.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4240.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4826.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3403.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.153.clone.1, %broadcast.4240.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1289.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2533.clone.1, %add.3403.clone.1), metadata={op_name="multiply.294"} - %div.2530.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3404.clone.1, %multiply.1289.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4677.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1367.clone.1, %broadcast.4242.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3402.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2530.clone.1, %mul.4677.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4676.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4678.clone.1, %add.3402.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3401.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1367.clone.1, %mul.4676.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.579 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3401.clone.1, %add.3401.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5052 = f32[]{:T(128)} constant(0) - %reduce.677 = f32[]{:T(128)} reduce(%square.579, %constant.5052), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.831.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3405.clone.1) - %bitcast.804.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3404.clone.1) - %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5052), dimensions={0,1,2,3}, to_apply=%region_191.216, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.664 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.677, %add.3401.clone.1, %bitcast.831.clone.1, %bitcast.804.clone.1, %reduce.686.clone.1) + %div.2545.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2199), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2544.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1367.clone.1, %div.2545.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2150.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2151.clone.1, %bitcast.1367.clone.1, %div.2544.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4843.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4269.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4843.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} + %mul.5008.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2150.clone.1, %broadcast.4269.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2945 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1366.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2945), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %constant.4842.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4268.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4842.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.5007.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4268.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3418.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.5008.clone.1, %mul.5007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4292 = f32[]{:T(128)S(6)} parameter(2) + %div.2543.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4292), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2150.clone.1, %select_n.2150.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4841.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4271.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} + %mul.5010.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4271.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5017 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1368.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5017), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %constant.4840.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4270.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.5009.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4270.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3419.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.5010.clone.1, %mul.5009.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4137 = f32[]{:T(128)S(6)} parameter(0) + %div.2542.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4137), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2541.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3419.clone.1, %div.2542.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.153.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2541.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4844.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4267.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4844.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %add.3417.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.153.clone.1, %broadcast.4267.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1289.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2543.clone.1, %add.3417.clone.1), metadata={op_name="multiply.294"} + %div.2540.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3418.clone.1, %multiply.1289.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.5005.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1365.clone.1, %broadcast.4269.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3416.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2540.clone.1, %mul.5005.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.5004.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.5006.clone.1, %add.3416.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3415.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1365.clone.1, %mul.5004.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.333 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3415.clone.1, %add.3415.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5070 = f32[]{:T(128)} constant(0) + %reduce.677 = f32[]{:T(128)} reduce(%square.333, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.829.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3419.clone.1) + %bitcast.802.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3418.clone.1) + %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_191.216, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.668 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.677, %add.3415.clone.1, %bitcast.829.clone.1, %bitcast.802.clone.1, %reduce.686.clone.1) } %region_155.180 (reduce_sum.438: f32[], reduce_sum.259: f32[]) -> f32[] { @@ -1620,62 +1620,62 @@ StackFrames ROOT %reduce_sum.260 = f32[]{:T(128)} add(%reduce_sum.438, %reduce_sum.259), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.529.clone.clone.clone (param_0.4079: bf16[4,128,129280], param_1.4953: s32[4,128], param_2.4225: f32[4,128], param_3.2913: f32[4,128], param_4.2170: bf16[4,128], param_5.1978: f32[4,128]) -> bf16[4,128,129280] { +%fused_computation.528.clone.clone.clone (param_0.4080: bf16[4,128,129280], param_1.4949: s32[4,128], param_2.4223: f32[4,128], param_3.2911: f32[4,128], param_4.2170: bf16[4,128], param_5.1978: f32[4,128]) -> bf16[4,128,129280] { %param_5.1978 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.4891 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.2913 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.4890 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2913), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.4079 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.3155 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4079), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %mul.5219 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.2911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.5218 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.4080 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.3161 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4080), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} %param_4.2170 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.791 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_4.2170), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.790 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.3155, %sub.791), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.534 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.790), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.4889 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4890, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.4225 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.2688 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4225), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.2687 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.4889, %div.2688), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.4953 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.363 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.4953), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.362 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.361 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.363, %eq.362), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.3154 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%eq.361), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.789 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%div.2687, %convert_element_type.3154), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.4888 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4891, %sub.789), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.3153 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.4888), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.939.clone.clone (param_0.4080: f32[4,128], param_1.4954: bf16[4,128,512], param_2.4227: bf16[512]) -> bf16[4,128,512] { - %param_2.4227 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) - %dot_general.831 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4227), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.4954 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3157 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4954), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.4080 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.4893 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4080), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.4892 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3157, %mul.4893), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.3156 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4892), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.830 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.831, %convert_element_type.3156), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.518 (param_0.4169: bf16[4,128,129280], param_1.5049: s32[4,128], param_2.4319: f32[4,128], param_3.2969: f32[4,128], param_4.2219: bf16[4,128], param_5.2020: f32[4,128], param_6.1457: f32[4,128], param_7.1138: bf16[4,128,512], param_8.902: bf16[512]) -> (f32[], bf16[512,129280,1]) { + %sub.804 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_4.2170), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.803 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.3161, %sub.804), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.534 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.803), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %mul.5217 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.5218, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.4223 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.2698 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4223), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.2697 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.5217, %div.2698), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.4949 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.371 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.4949), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.370 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.369 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.371, %eq.370), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %convert_element_type.3160 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%eq.369), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.802 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%div.2697, %convert_element_type.3160), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.5216 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.5219, %sub.802), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.3159 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.5216), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.938.clone.clone (param_0.4081: f32[4,128], param_1.4950: bf16[4,128,512], param_2.4225: bf16[512]) -> bf16[4,128,512] { + %param_2.4225 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.831 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4225), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.4950 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.3163 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.4081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.5221 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4081), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.5220 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3163, %mul.5221), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.3162 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.5220), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.830 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.831, %convert_element_type.3162), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.517 (param_0.4170: bf16[4,128,129280], param_1.5045: s32[4,128], param_2.4317: f32[4,128], param_3.2967: f32[4,128], param_4.2219: bf16[4,128], param_5.2020: f32[4,128], param_6.1457: f32[4,128], param_7.1138: bf16[4,128,512], param_8.902: bf16[512]) -> (f32[], bf16[512,129280,1]) { %param_6.1457 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) %param_7.1138 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) %param_8.902 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(8) - %fusion.577.clone.1 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_6.1457, %param_7.1138, %param_8.902), kind=kLoop, calls=%fused_computation.939.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.4169 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.5049 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4319 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.2969 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %fusion.574.clone.1 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_6.1457, %param_7.1138, %param_8.902), kind=kLoop, calls=%fused_computation.938.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.4170 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.5045 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.4317 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.2967 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) %param_4.2219 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) %param_5.2020 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} fusion(%param_0.4169, %param_1.5049, %param_2.4319, %param_3.2969, %param_4.2219, /*index=5*/%param_5.2020), kind=kLoop, calls=%fused_computation.529.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.141.clone.1 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.577.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.776 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%convolution.141.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.2657 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.776), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %square.581 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2657, %convert_element_type.2657), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5085 = f32[]{:T(128)} constant(0) - %reduce.678 = f32[]{:T(128)} reduce(%square.581, %constant.5085), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.753 = (f32[]{:T(128)}, bf16[512,129280,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.678, %convolution.141.clone.1) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} fusion(%param_0.4170, %param_1.5045, %param_2.4317, %param_3.2967, %param_4.2219, /*index=5*/%param_5.2020), kind=kLoop, calls=%fused_computation.528.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convolution.141.clone.1 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.574.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %bitcast.774 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%convolution.141.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.2663 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.774), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %mul.3871 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2663, %convert_element_type.2663), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5103 = f32[]{:T(128)} constant(0) + %reduce.678 = f32[]{:T(128)} reduce(%mul.3871, %constant.5103), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.757 = (f32[]{:T(128)}, bf16[512,129280,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.678, %convolution.141.clone.1) } %region_174.199 (reduce_sum.564: f32[], reduce_sum.387: f32[]) -> f32[] { @@ -1684,12 +1684,12 @@ StackFrames ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.564, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.519 (param_0.4153: bf16[129280,512]) -> f32[] { - %param_0.4153 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2659 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4153), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %square.583 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2659, %convert_element_type.2659), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5069 = f32[]{:T(128)} constant(0) - ROOT %reduce.679 = f32[]{:T(128)} reduce(%square.583, %constant.5069), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.518 (param_0.4154: bf16[129280,512]) -> f32[] { + %param_0.4154 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2665 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4154), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %mul.3873 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2665, %convert_element_type.2665), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5087 = f32[]{:T(128)} constant(0) + ROOT %reduce.679 = f32[]{:T(128)} reduce(%mul.3873, %constant.5087), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_240.265 (reduce_sum.1026: f32[], reduce_sum.689: f32[]) -> f32[] { @@ -1704,55 +1704,55 @@ StackFrames ROOT %reduce_sum.534 = f32[]{:T(128)} add(%reduce_sum.788, %reduce_sum.533), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.520 (param_0.4121: f32[129280,512], param_1.5006: f32[], param_2.4279: f32[], param_3.2932: f32[], param_4.2184: f32[129280,512], param_5.1987: f32[], param_6.1424: bf16[129280,512], param_7.1105: pred[], param_8.870: f32[129280,512]) -> (f32[], f32[129280,512], f32[129280,512], f32[129280,512], f32[]) { - %param_0.4121 = f32[129280,512]{1,0:T(8,128)} parameter(0) - %param_3.2932 = f32[]{:T(128)S(6)} parameter(3) - %mul.4552.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2932), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.519 (param_0.4122: f32[129280,512], param_1.5002: f32[], param_2.4277: f32[], param_3.2930: f32[], param_4.2184: f32[129280,512], param_5.1987: f32[], param_6.1424: bf16[129280,512], param_7.1105: pred[], param_8.870: f32[129280,512]) -> (f32[], f32[129280,512], f32[129280,512], f32[129280,512], f32[]) { + %param_0.4122 = f32[129280,512]{1,0:T(8,128)} parameter(0) + %param_3.2930 = f32[]{:T(128)S(6)} parameter(3) + %mul.4880.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1105 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2061.clone.1 = pred[129280,512]{1,0:T(8,128)(4,1)} broadcast(%param_7.1105), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2105.clone.1 = pred[129280,512]{1,0:T(8,128)(4,1)} broadcast(%param_7.1105), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1424 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.3098.clone.1 = f32[129280,512]{1,0:T(8,128)} convert(%param_6.1424), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convert_element_type.3104.clone.1 = f32[129280,512]{1,0:T(8,128)} convert(%param_6.1424), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %param_5.1987 = f32[]{:T(128)} parameter(5) - %div.2429.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_5.1987), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2428.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%convert_element_type.3098.clone.1, %div.2429.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2060.clone.1 = f32[129280,512]{1,0:T(8,128)} select(%select_n.2061.clone.1, %convert_element_type.3098.clone.1, %div.2428.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4735.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4182.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4735.clone.1), dimensions={}, metadata={op_name="broadcast.318"} - %mul.4558.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %broadcast.4182.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2439.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_5.1987), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2438.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%convert_element_type.3104.clone.1, %div.2439.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2104.clone.1 = f32[129280,512]{1,0:T(8,128)} select(%select_n.2105.clone.1, %convert_element_type.3104.clone.1, %div.2438.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4753.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4209.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4753.clone.1), dimensions={}, metadata={op_name="broadcast.318"} + %mul.4886.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4209.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.870 = f32[129280,512]{1,0:T(8,128)} parameter(8) - %constant.4739.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4559.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4739.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4557.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3324.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4558.clone.1, %mul.4557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4279 = f32[]{:T(128)S(6)} parameter(2) - %div.2425.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4279), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.380.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %select_n.2060.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4738.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4556.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4738.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4554.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4757.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4887.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4757.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4885.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4887.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3338.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4886.clone.1, %mul.4885.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4277 = f32[]{:T(128)S(6)} parameter(2) + %div.2435.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.380.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2104.clone.1, %select_n.2104.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4756.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4884.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4756.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4882.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2184 = f32[129280,512]{1,0:T(8,128)} parameter(4) - %constant.4737.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4555.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4737.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4553.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4555.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3323.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4554.clone.1, %mul.4553.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5006 = f32[]{:T(128)S(6)} parameter(1) - %div.2424.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2423.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3323.clone.1, %div.2424.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.138.clone.1 = f32[129280,512]{1,0:T(8,128)} sqrt(%div.2423.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4736.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3322.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4736.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3321.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%sqrt.138.clone.1, %add.3322.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1274.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%div.2425.clone.1, %add.3321.clone.1), metadata={op_name="multiply.309"} - %div.2422.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3324.clone.1, %multiply.1274.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4551.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4121, %broadcast.4182.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3320.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2422.clone.1, %mul.4551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4550.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4552.clone.1, %add.3320.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3319.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4121, %mul.4550.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.584 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3319.clone.1, %add.3319.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5037 = f32[]{:T(128)} constant(0) - %reduce.680 = f32[]{:T(128)} reduce(%square.584, %constant.5037), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.687.clone.1 = f32[]{:T(128)} reduce(%integer_pow.380.clone.1, %constant.5037), dimensions={0,1}, to_apply=%region_206.231, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.667 = (f32[]{:T(128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.680, %add.3319.clone.1, %add.3323.clone.1, %add.3324.clone.1, %reduce.687.clone.1) + %constant.4755.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4883.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4755.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4881.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3337.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4882.clone.1, %mul.4881.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5002 = f32[]{:T(128)S(6)} parameter(1) + %div.2434.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5002), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2433.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3337.clone.1, %div.2434.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.138.clone.1 = f32[129280,512]{1,0:T(8,128)} sqrt(%div.2433.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4754.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3336.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4754.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3335.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%sqrt.138.clone.1, %add.3336.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1274.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%div.2435.clone.1, %add.3335.clone.1), metadata={op_name="multiply.309"} + %div.2432.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3338.clone.1, %multiply.1274.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4879.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4122, %broadcast.4209.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3334.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2432.clone.1, %mul.4879.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4878.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4880.clone.1, %add.3334.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3333.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4122, %mul.4878.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.334 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3333.clone.1, %add.3333.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5055 = f32[]{:T(128)} constant(0) + %reduce.680 = f32[]{:T(128)} reduce(%square.334, %constant.5055), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.687.clone.1 = f32[]{:T(128)} reduce(%integer_pow.380.clone.1, %constant.5055), dimensions={0,1}, to_apply=%region_206.231, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.671 = (f32[]{:T(128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.680, %add.3333.clone.1, %add.3337.clone.1, %add.3338.clone.1, %reduce.687.clone.1) } %region_222.247 (reduce_sum.900: f32[], reduce_sum.605: f32[]) -> f32[] { @@ -1767,56 +1767,56 @@ StackFrames ROOT %reduce_sum.455 = f32[]{:T(128)} add(%reduce_sum.662, %reduce_sum.451), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.521 (param_0.4139: f32[512,129280], param_1.5024: f32[], param_2.4297: f32[], param_3.2950: f32[], param_4.2202: f32[512,129280], param_5.2005: f32[], param_6.1442: bf16[512,129280,1], param_7.1123: pred[], param_8.888: f32[512,129280]) -> (f32[], f32[512,129280], f32[512,129280], f32[512,129280], f32[]) { - %param_0.4139 = f32[512,129280]{1,0:T(8,128)} parameter(0) - %param_3.2950 = f32[]{:T(128)S(6)} parameter(3) - %mul.4705.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2950), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.520 (param_0.4140: f32[512,129280], param_1.5020: f32[], param_2.4295: f32[], param_3.2948: f32[], param_4.2202: f32[512,129280], param_5.2005: f32[], param_6.1442: bf16[512,129280,1], param_7.1123: pred[], param_8.888: f32[512,129280]) -> (f32[], f32[512,129280], f32[512,129280], f32[512,129280], f32[]) { + %param_0.4140 = f32[512,129280]{1,0:T(8,128)} parameter(0) + %param_3.2948 = f32[]{:T(128)S(6)} parameter(3) + %mul.5033.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2948), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1123 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2117.clone.1 = pred[512,129280]{1,0:T(8,128)(4,1)} broadcast(%param_7.1123), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2161.clone.1 = pred[512,129280]{1,0:T(8,128)(4,1)} broadcast(%param_7.1123), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1442 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.1372.clone.1 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%param_6.1442), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.3100.clone.1 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.1372.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %bitcast.1370.clone.1 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%param_6.1442), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.3106.clone.1 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.1370.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} %param_5.2005 = f32[]{:T(128)} parameter(5) - %div.2557.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_5.2005), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2556.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%convert_element_type.3100.clone.1, %div.2557.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2116.clone.1 = f32[512,129280]{1,0:T(8,128)} select(%select_n.2117.clone.1, %convert_element_type.3100.clone.1, %div.2556.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4839.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4250.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4839.clone.1), dimensions={}, metadata={op_name="broadcast.333"} - %mul.4711.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2567.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_5.2005), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2566.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%convert_element_type.3106.clone.1, %div.2567.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2160.clone.1 = f32[512,129280]{1,0:T(8,128)} select(%select_n.2161.clone.1, %convert_element_type.3106.clone.1, %div.2566.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4857.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4277.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4857.clone.1), dimensions={}, metadata={op_name="broadcast.333"} + %mul.5039.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2160.clone.1, %broadcast.4277.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.888 = f32[512,129280]{1,0:T(8,128)} parameter(8) - %constant.4843.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4712.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4843.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4710.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.4712.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3423.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4711.clone.1, %mul.4710.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4297 = f32[]{:T(128)S(6)} parameter(2) - %div.2553.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4297), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.398.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %select_n.2116.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4842.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4709.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4842.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4707.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.4709.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4861.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.5040.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4861.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5038.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.5040.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3437.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.5039.clone.1, %mul.5038.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4295 = f32[]{:T(128)S(6)} parameter(2) + %div.2563.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4295), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.398.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2160.clone.1, %select_n.2160.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4860.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.5037.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5035.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.5037.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2202 = f32[512,129280]{1,0:T(8,128)} parameter(4) - %constant.4841.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4708.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4706.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.4708.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3422.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4707.clone.1, %mul.4706.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5024 = f32[]{:T(128)S(6)} parameter(1) - %div.2552.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5024), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2551.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3422.clone.1, %div.2552.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.156.clone.1 = f32[512,129280]{1,0:T(8,128)} sqrt(%div.2551.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4840.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3421.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3420.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%sqrt.156.clone.1, %add.3421.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1292.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%div.2553.clone.1, %add.3420.clone.1), metadata={op_name="multiply.291"} - %div.2550.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3423.clone.1, %multiply.1292.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4704.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4139, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3419.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2550.clone.1, %mul.4704.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4703.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.4705.clone.1, %add.3419.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3418.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4139, %mul.4703.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.585 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3418.clone.1, %add.3418.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5055 = f32[]{:T(128)} constant(0) - %reduce.681 = f32[]{:T(128)} reduce(%square.585, %constant.5055), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.688.clone.1 = f32[]{:T(128)} reduce(%integer_pow.398.clone.1, %constant.5055), dimensions={0,1}, to_apply=%region_188.213, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.668 = (f32[]{:T(128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.681, %add.3418.clone.1, %add.3422.clone.1, %add.3423.clone.1, %reduce.688.clone.1) + %constant.4859.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.5036.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4859.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5034.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.5036.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3436.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.5035.clone.1, %mul.5034.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5020 = f32[]{:T(128)S(6)} parameter(1) + %div.2562.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5020), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2561.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3436.clone.1, %div.2562.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.156.clone.1 = f32[512,129280]{1,0:T(8,128)} sqrt(%div.2561.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4858.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3435.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3434.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%sqrt.156.clone.1, %add.3435.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1292.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%div.2563.clone.1, %add.3434.clone.1), metadata={op_name="multiply.291"} + %div.2560.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3437.clone.1, %multiply.1292.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.5032.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4140, %broadcast.4277.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3433.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2560.clone.1, %mul.5032.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.5031.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.5033.clone.1, %add.3433.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3432.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4140, %mul.5031.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.335 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3432.clone.1, %add.3432.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5073 = f32[]{:T(128)} constant(0) + %reduce.681 = f32[]{:T(128)} reduce(%square.335, %constant.5073), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.688.clone.1 = f32[]{:T(128)} reduce(%integer_pow.398.clone.1, %constant.5073), dimensions={0,1}, to_apply=%region_188.213, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.672 = (f32[]{:T(128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.681, %add.3432.clone.1, %add.3436.clone.1, %add.3437.clone.1, %reduce.688.clone.1) } %region_207.232 (reduce_sum.795: f32[], reduce_sum.535: f32[]) -> f32[] { @@ -1825,23 +1825,23 @@ StackFrames ROOT %reduce_sum.540 = f32[]{:T(128)} add(%reduce_sum.795, %reduce_sum.535), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.522 (param_0.4190: bf16[4,128,129280], param_1.5063: f32[4,128], param_2.4329: s32[4,128], param_3.2977: bf16[4,128]) -> f32[4,128] { - %param_2.4329 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.299 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4329), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.294 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.293 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.299, %eq.294), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.4190 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2664 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4190), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.2977 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.652 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2977), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.643 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2664, %sub.652), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.5063 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.650 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5063), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.639 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%sub.643, %sub.650), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.5109 = f32[]{:T(128)} constant(0) - %broadcast.3757 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5109), dimensions={}, metadata={op_name="broadcast.514"} - %mul.3612 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.293, %sub.639, %broadcast.3757), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3612, %constant.5109), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +%fused_computation.521 (param_0.4191: bf16[4,128,129280], param_1.5059: f32[4,128], param_2.4327: s32[4,128], param_3.2975: bf16[4,128]) -> f32[4,128] { + %param_2.4327 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.307 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4327), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.302 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.301 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.307, %eq.302), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %param_0.4191 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2670 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_3.2975 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) + %sub.665 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2975), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.656 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2670, %sub.665), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.5059 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.663 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5059), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.652 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%sub.656, %sub.663), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %constant.5127 = f32[]{:T(128)} constant(0) + %broadcast.3784 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5127), dimensions={}, metadata={op_name="broadcast.518"} + %mul.3874 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.301, %sub.652, %broadcast.3784), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3874, %constant.5127), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_37.47 (reduce_sum.76: f32[], reduce_sum.80: f32[]) -> f32[] { @@ -1850,15 +1850,15 @@ StackFrames ROOT %reduce_sum.83 = f32[]{:T(128)} add(%reduce_sum.76, %reduce_sum.80), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.533 (param_0.4191: bf16[4,128,129280], param_1.5064: bf16[4,128]) -> f32[4,128] { - %param_0.4191 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2670 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.5064 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.653 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5064), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.649 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2670, %sub.653), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.448 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.649), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.5110 = f32[]{:T(128)} constant(0) - ROOT %reduce.683 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.448, %constant.5110), dimensions={2}, to_apply=%region_37.47, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +%fused_computation.532 (param_0.4192: bf16[4,128,129280], param_1.5060: bf16[4,128]) -> f32[4,128] { + %param_0.4192 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2676 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.5060 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.666 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5060), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.662 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2676, %sub.666), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.448 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.662), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %constant.5128 = f32[]{:T(128)} constant(0) + ROOT %reduce.683 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.448, %constant.5128), dimensions={2}, to_apply=%region_37.47, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_152.177 (reduce_sum.417: f32[], reduce_sum.244: f32[]) -> f32[] { @@ -1867,18 +1867,18 @@ StackFrames ROOT %reduce_sum.251 = f32[]{:T(128)} add(%reduce_sum.417, %reduce_sum.244), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.541 (param_0.4172: f32[3,512,128,256]) -> f32[] { - %param_0.4172 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.752 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4172), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.588 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.752, %bitcast.752), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5088 = f32[]{:T(128)} constant(0) - ROOT %reduce.689 = f32[]{:T(128)} reduce(%square.588, %constant.5088), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.540 (param_0.4173: f32[3,512,128,256]) -> f32[] { + %param_0.4173 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.750 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4173), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3895 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.750, %bitcast.750), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5106 = f32[]{:T(128)} constant(0) + ROOT %reduce.689 = f32[]{:T(128)} reduce(%mul.3895, %constant.5106), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.542 (param_0.1602: f32[512,3,128,256]) -> bf16[3,512,128,256] { - %param_0.1602 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) - %copy.1551 = bf16[512,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1602), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\']"} - ROOT %bitcast.753 = bf16[3,512,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1551), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} +%fused_computation.541 (param_0.1600: f32[512,3,128,256]) -> bf16[3,512,128,256] { + %param_0.1600 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) + %copy.1551 = bf16[512,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1600), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\']"} + ROOT %bitcast.751 = bf16[3,512,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1551), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} } %region_219.244 (reduce_sum.879: f32[], reduce_sum.591: f32[]) -> f32[] { @@ -1893,55 +1893,55 @@ StackFrames ROOT %reduce_sum.442 = f32[]{:T(128)} add(%reduce_sum.641, %reduce_sum.437), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.543 (param_0.4142: f32[512,3,128,256], param_1.5027: f32[], param_2.4300: f32[], param_3.2953: f32[], param_4.2205: f32[512,3,128,256], param_5.2008: f32[], param_6.1445: f32[3,512,128,256], param_7.1126: pred[], param_8.891: f32[512,3,128,256]) -> (f32[], f32[512,3,128,256], f32[512,3,128,256], f32[512,3,128,256], f32[]) { - %param_0.4142 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) - %param_3.2953 = f32[]{:T(128)S(6)} parameter(3) - %mul.4735.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2953), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.542 (param_0.4143: f32[512,3,128,256], param_1.5023: f32[], param_2.4298: f32[], param_3.2951: f32[], param_4.2205: f32[512,3,128,256], param_5.2008: f32[], param_6.1445: f32[3,512,128,256], param_7.1126: pred[], param_8.891: f32[512,3,128,256]) -> (f32[], f32[512,3,128,256], f32[512,3,128,256], f32[512,3,128,256], f32[]) { + %param_0.4143 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) + %param_3.2951 = f32[]{:T(128)S(6)} parameter(3) + %mul.5063.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1126 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2129.clone.1 = pred[512,3,128,256]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.1126), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2173.clone.1 = pred[512,3,128,256]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.1126), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1445 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.1378.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_6.1445), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1376.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_6.1445), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} %param_5.2008 = f32[]{:T(128)} parameter(5) - %div.2581.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_5.2008), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2580.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.1378.clone.1, %div.2581.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2128.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2129.clone.1, %bitcast.1378.clone.1, %div.2580.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4857.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4256.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4857.clone.1), dimensions={}, metadata={op_name="broadcast.336"} - %mul.4741.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2591.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_5.2008), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2590.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.1376.clone.1, %div.2591.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2172.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2173.clone.1, %bitcast.1376.clone.1, %div.2590.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4875.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4283.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4875.clone.1), dimensions={}, metadata={op_name="broadcast.336"} + %mul.5069.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2172.clone.1, %broadcast.4283.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.891 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(8) - %constant.4861.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4742.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4861.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4740.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.4742.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3441.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4741.clone.1, %mul.4740.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4300 = f32[]{:T(128)S(6)} parameter(2) - %div.2577.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4300), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.401.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %select_n.2128.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4860.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4739.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4737.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.4739.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4879.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.5070.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4879.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5068.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.5070.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3455.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.5069.clone.1, %mul.5068.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4298 = f32[]{:T(128)S(6)} parameter(2) + %div.2587.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4298), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.401.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2172.clone.1, %select_n.2172.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4878.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.5067.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4878.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5065.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.5067.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2205 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(4) - %constant.4859.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4738.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4859.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4736.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.4738.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3440.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4737.clone.1, %mul.4736.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5027 = f32[]{:T(128)S(6)} parameter(1) - %div.2576.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5027), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2575.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3440.clone.1, %div.2576.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.159.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} sqrt(%div.2575.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4858.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3439.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3438.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3439.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1295.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%div.2577.clone.1, %add.3438.clone.1), metadata={op_name="multiply.288"} - %div.2574.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3441.clone.1, %multiply.1295.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4734.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4142, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3437.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2574.clone.1, %mul.4734.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4733.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.4735.clone.1, %add.3437.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3436.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4142, %mul.4733.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.589 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3436.clone.1, %add.3436.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5058 = f32[]{:T(128)} constant(0) - %reduce.690 = f32[]{:T(128)} reduce(%square.589, %constant.5058), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.691.clone.1 = f32[]{:T(128)} reduce(%integer_pow.401.clone.1, %constant.5058), dimensions={0,1,2,3}, to_apply=%region_185.210, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.663 = (f32[]{:T(128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.690, %add.3436.clone.1, %add.3440.clone.1, %add.3441.clone.1, %reduce.691.clone.1) + %constant.4877.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.5066.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4877.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5064.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.5066.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3454.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.5065.clone.1, %mul.5064.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5023 = f32[]{:T(128)S(6)} parameter(1) + %div.2586.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5023), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2585.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3454.clone.1, %div.2586.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.159.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} sqrt(%div.2585.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4876.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3453.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4876.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3452.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3453.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1295.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%div.2587.clone.1, %add.3452.clone.1), metadata={op_name="multiply.288"} + %div.2584.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3455.clone.1, %multiply.1295.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.5062.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4143, %broadcast.4283.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3451.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2584.clone.1, %mul.5062.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.5061.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.5063.clone.1, %add.3451.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3450.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4143, %mul.5061.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.336 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3450.clone.1, %add.3450.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5076 = f32[]{:T(128)} constant(0) + %reduce.690 = f32[]{:T(128)} reduce(%square.336, %constant.5076), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.691.clone.1 = f32[]{:T(128)} reduce(%integer_pow.401.clone.1, %constant.5076), dimensions={0,1,2,3}, to_apply=%region_185.210, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.667 = (f32[]{:T(128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.690, %add.3450.clone.1, %add.3454.clone.1, %add.3455.clone.1, %reduce.691.clone.1) } %region_172.197 (reduce_sum.557: f32[], reduce_sum.381: f32[]) -> f32[] { @@ -1950,39 +1950,39 @@ StackFrames ROOT %reduce_sum.386 = f32[]{:T(128)} add(%reduce_sum.557, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.783.clone.clone (param_0.4106: f32[4,128], param_1.4998: bf16[4,128,1536], param_2.4261: bf16[1536]) -> bf16[4,128,1536,1] { - %param_2.4261 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.851 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4261), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.4998 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3179 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4998), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} - %param_0.4106 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.4939 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4106), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %mul.4938 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3179, %mul.4939), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %convert_element_type.3178 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.4938), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} - %dot_general.850 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.851, %convert_element_type.3178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.1466 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dot_general.850), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.782.clone.clone (param_0.4107: f32[4,128], param_1.4994: bf16[4,128,1536], param_2.4259: bf16[1536]) -> bf16[4,128,1536,1] { + %param_2.4259 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.851 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4259), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.4994 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.3185 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4994), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} + %param_0.4107 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.5267 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4107), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} + %mul.5266 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3185, %mul.5267), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} + %convert_element_type.3184 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.5266), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} + %dot_general.850 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.851, %convert_element_type.3184), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.1464 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dot_general.850), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} } %bitcast_fusion.12 (bitcast_input.12: bf16[4,128,128,192]) -> bf16[4,128,128,192] { %bitcast_input.12 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.1488 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} bitcast(%bitcast_input.12) -} - -%fused_computation.552 (param_0.4154: bf16[4,128,128,192], param_1.5038: f32[4,128], param_2.4311: bf16[4,128,1536], param_3.2964: bf16[1536]) -> (f32[], bf16[1536,128,192,1]) { - %param_1.5038 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4311 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.2964 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.460.clone.1 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.5038, %param_2.4311, %param_3.2964), kind=kLoop, calls=%fused_computation.783.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.4154 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - %fusion.751 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.4154), kind=kLoop, calls=%bitcast_fusion.12 - %convolution.146.clone.1 = bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)} convolution(%fusion.460.clone.1, %fusion.751), window={size=192x4 pad=191_191x0_0 rhs_reversal=1x0}, dim_labels=1fb0_1io0->bf01, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} - %bitcast.861 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%convolution.146.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} - %broadcast_in_dim.1275 = f32[1536,128,192]{1,0,2:T(8,128)} convert(%bitcast.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.763 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1275), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.592 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.763, %bitcast.763), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5070 = f32[]{:T(128)} constant(0) - %reduce.692 = f32[]{:T(128)} reduce(%square.592, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.762 = (f32[]{:T(128)}, bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)}) tuple(%reduce.692, %convolution.146.clone.1) + ROOT %bitcast.1486 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} bitcast(%bitcast_input.12) +} + +%fused_computation.551 (param_0.4155: bf16[4,128,128,192], param_1.5034: f32[4,128], param_2.4309: bf16[4,128,1536], param_3.2962: bf16[1536]) -> (f32[], bf16[1536,128,192,1]) { + %param_1.5034 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.4309 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.2962 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.457.clone.1 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.5034, %param_2.4309, %param_3.2962), kind=kLoop, calls=%fused_computation.782.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.4155 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) + %fusion.748 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.4155), kind=kLoop, calls=%bitcast_fusion.12 + %convolution.144.clone.1 = bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)} convolution(%fusion.457.clone.1, %fusion.748), window={size=192x4 pad=191_191x0_0 rhs_reversal=1x0}, dim_labels=1fb0_1io0->bf01, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} + %bitcast.859 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%convolution.144.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} + %broadcast_in_dim.1388 = f32[1536,128,192]{1,0,2:T(8,128)} convert(%bitcast.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.761 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3904 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.761, %bitcast.761), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5088 = f32[]{:T(128)} constant(0) + %reduce.692 = f32[]{:T(128)} reduce(%mul.3904, %constant.5088), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.766 = (f32[]{:T(128)}, bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)}) tuple(%reduce.692, %convolution.144.clone.1) } %region_239.264 (reduce_sum.1019: f32[], reduce_sum.687: f32[]) -> f32[] { @@ -1997,4 +1997,4 @@ StackFrames ROOT %reduce_sum.528 = f32[]{:T(128)} add(%reduce_sum.781, %reduce_sum.527), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.557 (param_0.4122: f32[], param_1.5007: f32[], param_2.4280: f32[], param_3.2933: f32[1536,1,128,192], param_4.2185: f32[1536,1,128,192], param_5.1988: f32[], param_6.1425: bf16[1536,128,192,1], param_7.1106: pred[], param_8.871: f32[1536,1,128,192]) -> (f32[], f32[1536,1,128,192], f32[1536,1,128,192], f32[1536,1,128,192], f32[]) { +%fused_computation.556 (param_0.4123: f32[], param_1.5003: f32[], param_2.4278: f32[], param_3.2931: f32[1536,1,128,192], param_4.2185: f32[1536,1,128,192], param_5.1988: f32[], param_6.1425: bf16[1536,128,192,1], param_7.1106: pred[], param_8.871: f32[1536,1,128,192]) -> (f32[], f32[1536,1,128,192], f32[1536,1,128,192], f32[1536,1,128,192], f32[]) { diff --git a/tests/utils/reference_hlo_llama3_8b.txt b/tests/utils/reference_hlo_llama3_8b.txt index 27c6529df2..2a4c292f01 100644 --- a/tests/utils/reference_hlo_llama3_8b.txt +++ b/tests/utils/reference_hlo_llama3_8b.txt @@ -44,62 +44,62 @@ StackFrames ROOT %reduce_sum.192 = f32[]{:T(128)} add(%reduce_sum.190, %reduce_sum.191), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.280.clone.clone.clone (param_0.1099: bf16[4,128,128256], param_1.1265: s32[4,128], param_2.1086: f32[4,128], param_3.785: f32[4,128], param_4.487: bf16[4,128], param_5.412: f32[4,128]) -> bf16[4,128,128256] { - %param_5.412 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1613 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.785 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1612 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.785), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1099 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1044 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1099), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.487 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.487), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1044, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.281.clone.clone.clone (param_0.1085: bf16[4,128,128256], param_1.1251: s32[4,128], param_2.1077: f32[4,128], param_3.781: f32[4,128], param_4.482: bf16[4,128], param_5.404: f32[4,128]) -> bf16[4,128,128256] { + %param_5.404 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.1679 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.404), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.781 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.1678 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.781), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1085 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1032 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1085), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.482 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.482), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1032, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.62 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1611 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1612, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1611, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1265 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1265), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.1677 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1678, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1077 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1077), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1677, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1251 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1251), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.48 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.47 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1043 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1043), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1610 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1613, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1042 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1610), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.316.clone.clone (param_0.1100: f32[4,128], param_1.1266: bf16[4,128,4096], param_2.1088: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1088 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.387 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1088), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1266 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1046 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1266), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1100 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1615 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1100), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1614 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1046, %mul.1615), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1045 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1614), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.387, %convert_element_type.1045), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.219 (param_0.1119: bf16[4,128,128256], param_1.1281: s32[4,128], param_2.1112: f32[4,128], param_3.801: f32[4,128], param_4.502: bf16[4,128], param_5.427: f32[4,128], param_6.299: f32[4,128], param_7.198: bf16[4,128,4096], param_8.116: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { - %param_6.299 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.198 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) - %param_8.116 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) - %fusion.239.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.299, %param_7.198, %param_8.116), kind=kLoop, calls=%fused_computation.316.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1119 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1281 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1112 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.801 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %param_4.502 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %param_5.427 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1119, %param_1.1281, %param_2.1112, %param_3.801, %param_4.502, /*index=5*/%param_5.427), kind=kLoop, calls=%fused_computation.280.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.88.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.239.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.306 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.88.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.923 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.306), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %square.157 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.923, %convert_element_type.923), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %convert_element_type.1031 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1031), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1676 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1679, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1030 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1676), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.317.clone.clone (param_0.1086: f32[4,128], param_1.1252: bf16[4,128,4096], param_2.1079: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1079 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1079), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1034 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1252), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1681 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1680 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1034, %mul.1681), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1033 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1680), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.220 (param_0.1105: bf16[4,128,128256], param_1.1267: s32[4,128], param_2.1103: f32[4,128], param_3.797: f32[4,128], param_4.497: bf16[4,128], param_5.419: f32[4,128], param_6.287: f32[4,128], param_7.186: bf16[4,128,4096], param_8.112: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { + %param_6.287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.186 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) + %param_8.112 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) + %fusion.229.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.287, %param_7.186, %param_8.112), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1105 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1267 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1103 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.797 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %param_4.497 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %param_5.419 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1105, %param_1.1267, %param_2.1103, %param_3.797, %param_4.497, /*index=5*/%param_5.419), kind=kLoop, calls=%fused_computation.281.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convolution.82.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.229.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %bitcast.300 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.82.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.911 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.300), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %mul.1350 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.911, %convert_element_type.911), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1006 = f32[]{:T(128)} constant(0) - %reduce.118 = f32[]{:T(128)} reduce(%square.157, %constant.1006), dimensions={0,1}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.88.clone.1) + %reduce.118 = f32[]{:T(128)} reduce(%mul.1350, %constant.1006), dimensions={0,1}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.82.clone.1) } %region_34.39 (reduce_sum.196: f32[], reduce_sum.197: f32[]) -> f32[] { @@ -108,12 +108,12 @@ StackFrames ROOT %reduce_sum.198 = f32[]{:T(128)} add(%reduce_sum.196, %reduce_sum.197), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.220 (param_0.1118: bf16[128256,4096]) -> f32[] { - %param_0.1118 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.925 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1118), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %square.159 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.925, %convert_element_type.925), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.221 (param_0.1104: bf16[128256,4096]) -> f32[] { + %param_0.1104 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.913 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1104), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %mul.1352 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.913, %convert_element_type.913), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1005 = f32[]{:T(128)} constant(0) - ROOT %reduce.119 = f32[]{:T(128)} reduce(%square.159, %constant.1005), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.119 = f32[]{:T(128)} reduce(%mul.1352, %constant.1005), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_60.65 (reduce_sum.338: f32[], reduce_sum.339: f32[]) -> f32[] { @@ -128,39 +128,39 @@ StackFrames ROOT %reduce_sum.261 = f32[]{:T(128)} add(%reduce_sum.259, %reduce_sum.260), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.221 (param_0.1106: f32[128256,4096], param_1.1269: f32[], param_2.1100: f32[], param_3.789: f32[], param_4.490: f32[128256,4096], param_5.415: f32[], param_6.287: bf16[128256,4096], param_7.186: pred[], param_8.104: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { - %param_0.1106 = f32[128256,4096]{1,0:T(8,128)} parameter(0) - %param_3.789 = f32[]{:T(128)S(6)} parameter(3) - %mul.1482.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.186 = pred[]{:T(512)S(6)} parameter(7) - %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.186), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.287 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1017.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.287), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %param_5.415 = f32[]{:T(128)} parameter(5) - %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1017.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1017.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.222 (param_0.1092: f32[128256,4096], param_1.1255: f32[], param_2.1091: f32[], param_3.785: f32[], param_4.485: f32[128256,4096], param_5.407: f32[], param_6.275: bf16[128256,4096], param_7.174: pred[], param_8.100: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { + %param_0.1092 = f32[128256,4096]{1,0:T(8,128)} parameter(0) + %param_3.785 = f32[]{:T(128)S(6)} parameter(3) + %mul.1548.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.785), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.174 = pred[]{:T(512)S(6)} parameter(7) + %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.174), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.275 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1005.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.275), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_5.407 = f32[]{:T(128)} parameter(5) + %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.407), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1005.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1005.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.907.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.554.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.907.clone.1), dimensions={}, metadata={op_name="broadcast.61"} - %mul.1488.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.104 = f32[128256,4096]{1,0:T(8,128)} parameter(8) + %mul.1554.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.100 = f32[128256,4096]{1,0:T(8,128)} parameter(8) %constant.911.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1489.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1487.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.104, %mul.1489.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1488.clone.1, %mul.1487.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) - %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1555.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1553.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.100, %mul.1555.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1554.clone.1, %mul.1553.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1091 = f32[]{:T(128)S(6)} parameter(2) + %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1091), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.60.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %select_n.241.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.910.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1486.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1484.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1486.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.490 = f32[128256,4096]{1,0:T(8,128)} parameter(4) + %mul.1552.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1550.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1552.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.485 = f32[128256,4096]{1,0:T(8,128)} parameter(4) %constant.909.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1485.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1483.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.490, %mul.1485.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1484.clone.1, %mul.1483.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1269 = f32[]{:T(128)S(6)} parameter(1) - %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1269), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1551.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1549.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.485, %mul.1551.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1550.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1255 = f32[]{:T(128)S(6)} parameter(1) + %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1255), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.719.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.775.clone.1, %div.720.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.58.clone.1 = f32[128256,4096]{1,0:T(8,128)} sqrt(%div.719.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.908.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -168,13 +168,13 @@ StackFrames %add.773.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%sqrt.58.clone.1, %add.774.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.256.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%div.721.clone.1, %add.773.clone.1), metadata={op_name="multiply.42"} %div.718.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.776.clone.1, %multiply.256.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1481.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1106, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1481.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1480.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1482.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1106, %mul.1480.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.160 = f32[128256,4096]{1,0:T(8,128)} multiply(%add.771.clone.1, %add.771.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1547.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1092, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1547.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1546.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1548.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1092, %mul.1546.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.118 = f32[128256,4096]{1,0:T(8,128)} multiply(%add.771.clone.1, %add.771.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.993 = f32[]{:T(128)} constant(0) - %reduce.120 = f32[]{:T(128)} reduce(%square.160, %constant.993), dimensions={0,1}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.120 = f32[]{:T(128)} reduce(%square.118, %constant.993), dimensions={0,1}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.122.clone.1 = f32[]{:T(128)} reduce(%integer_pow.60.clone.1, %constant.993), dimensions={0,1}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.135 = (f32[]{:T(128)}, f32[128256,4096]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.120, %add.771.clone.1, %add.775.clone.1, %add.776.clone.1, %reduce.122.clone.1) } @@ -191,40 +191,40 @@ StackFrames ROOT %reduce_sum.255 = f32[]{:T(128)} add(%reduce_sum.253, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.222 (param_0.1107: f32[4096,128256], param_1.1270: f32[], param_2.1101: f32[], param_3.790: f32[], param_4.491: f32[4096,128256], param_5.416: f32[], param_6.288: bf16[4096,128256,1], param_7.187: pred[], param_8.105: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { - %param_0.1107 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %param_3.790 = f32[]{:T(128)S(6)} parameter(3) - %mul.1492.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.187 = pred[]{:T(512)S(6)} parameter(7) - %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.187), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.288 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.409.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.288), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.1019.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.409.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %param_5.416 = f32[]{:T(128)} parameter(5) - %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1019.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1019.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.223 (param_0.1093: f32[4096,128256], param_1.1256: f32[], param_2.1092: f32[], param_3.786: f32[], param_4.486: f32[4096,128256], param_5.408: f32[], param_6.276: bf16[4096,128256,1], param_7.175: pred[], param_8.101: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { + %param_0.1093 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %param_3.786 = f32[]{:T(128)S(6)} parameter(3) + %mul.1558.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.786), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.175 = pred[]{:T(512)S(6)} parameter(7) + %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.175), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.276 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) + %bitcast.403.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.1007.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.403.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %param_5.408 = f32[]{:T(128)} parameter(5) + %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.408), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1007.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1007.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.913.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.556.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.913.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %mul.1498.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.105 = f32[4096,128256]{1,0:T(8,128)} parameter(8) + %mul.1564.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.101 = f32[4096,128256]{1,0:T(8,128)} parameter(8) %constant.917.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1499.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1497.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.105, %mul.1499.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1498.clone.1, %mul.1497.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) - %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1565.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1563.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.101, %mul.1565.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1564.clone.1, %mul.1563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1092 = f32[]{:T(128)S(6)} parameter(2) + %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1092), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.61.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %select_n.245.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.916.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1496.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1494.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1496.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.491 = f32[4096,128256]{1,0:T(8,128)} parameter(4) + %mul.1562.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1560.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.486 = f32[4096,128256]{1,0:T(8,128)} parameter(4) %constant.915.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1495.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1493.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.491, %mul.1495.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1494.clone.1, %mul.1493.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1270 = f32[]{:T(128)S(6)} parameter(1) - %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1270), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1561.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1559.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.486, %mul.1561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1560.clone.1, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1256 = f32[]{:T(128)S(6)} parameter(1) + %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1256), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.727.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.781.clone.1, %div.728.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.59.clone.1 = f32[4096,128256]{1,0:T(8,128)} sqrt(%div.727.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.914.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -232,13 +232,13 @@ StackFrames %add.779.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%sqrt.59.clone.1, %add.780.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.257.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%div.729.clone.1, %add.779.clone.1), metadata={op_name="multiply.41"} %div.726.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.782.clone.1, %multiply.257.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1491.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1107, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1491.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1490.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1492.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1107, %mul.1490.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.161 = f32[4096,128256]{1,0:T(8,128)} multiply(%add.777.clone.1, %add.777.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1557.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1093, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1556.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1558.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1093, %mul.1556.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.119 = f32[4096,128256]{1,0:T(8,128)} multiply(%add.777.clone.1, %add.777.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.994 = f32[]{:T(128)} constant(0) - %reduce.121 = f32[]{:T(128)} reduce(%square.161, %constant.994), dimensions={0,1}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.121 = f32[]{:T(128)} reduce(%square.119, %constant.994), dimensions={0,1}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.123.clone.1 = f32[]{:T(128)} reduce(%integer_pow.61.clone.1, %constant.994), dimensions={0,1}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.136 = (f32[]{:T(128)}, f32[4096,128256]{1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.121, %add.777.clone.1, %add.781.clone.1, %add.782.clone.1, %reduce.123.clone.1) } @@ -249,12 +249,12 @@ StackFrames ROOT %reduce_sum.156 = f32[]{:T(128)} add(%reduce_sum.154, %reduce_sum.155), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.239 (param_0.1124: f32[4,14336,4096]) -> f32[] { - %param_0.1124 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) - %bitcast.314 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1124), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.164 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.314, %bitcast.314), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.240 (param_0.1110: f32[4,14336,4096]) -> f32[] { + %param_0.1110 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) + %bitcast.308 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1375 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.308, %bitcast.308), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1011 = f32[]{:T(128)} constant(0) - ROOT %reduce.124 = f32[]{:T(128)} reduce(%square.164, %constant.1011), dimensions={0,1,2}, to_apply=%region_25.30, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.124 = f32[]{:T(128)} reduce(%mul.1375, %constant.1011), dimensions={0,1,2}, to_apply=%region_25.30, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_24.29 (reduce_sum.148: f32[], reduce_sum.149: f32[]) -> f32[] { @@ -269,35 +269,35 @@ StackFrames ROOT %reduce_sum.147 = f32[]{:T(128)} add(%reduce_sum.142, %reduce_sum.143), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.241 (param_0.1125: f32[4,4096,14336], param_1.1284: f32[4,4096,14336]) -> (f32[], f32[]) { - %param_0.1125 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) - %bitcast.318 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1125), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.167 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.318, %bitcast.318), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.242 (param_0.1111: f32[4,4096,14336], param_1.1270: f32[4,4096,14336]) -> (f32[], f32[]) { + %param_0.1111 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) + %bitcast.312 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1378 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.312, %bitcast.312), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1012 = f32[]{:T(128)} constant(0) - %reduce.125 = f32[]{:T(128)} reduce(%square.167, %constant.1012), dimensions={0,1,2}, to_apply=%region_24.29, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) - %bitcast.322.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.170.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.322.clone.1, %bitcast.322.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.126.clone.1 = f32[]{:T(128)} reduce(%square.170.clone.1, %constant.1012), dimensions={0,1,2}, to_apply=%region_23.28, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.125 = f32[]{:T(128)} reduce(%mul.1378, %constant.1012), dimensions={0,1,2}, to_apply=%region_24.29, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1270 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) + %bitcast.316.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1270), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1381.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.316.clone.1, %bitcast.316.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.126.clone.1 = f32[]{:T(128)} reduce(%mul.1381.clone.1, %constant.1012), dimensions={0,1,2}, to_apply=%region_23.28, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.155 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.125, %reduce.126.clone.1) } -%fused_computation.244 (param_0.694: f32[14336,4,4096]) -> bf16[4,14336,4096] { - %param_0.694 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.694), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.323 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.245 (param_0.681: f32[14336,4,4096]) -> bf16[4,14336,4096] { + %param_0.681 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) + %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.681), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} + ROOT %bitcast.317 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.245 (param_0.696: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.696 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.696), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.324 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.246 (param_0.683: f32[4096,4,14336]) -> bf16[4,4096,14336] { + %param_0.683 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.683), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} + ROOT %bitcast.318 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.246 (param_0.698: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.698 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.698), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.325 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.247 (param_0.685: f32[4096,4,14336]) -> bf16[4,4096,14336] { + %param_0.685 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.685), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} + ROOT %bitcast.319 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_52.57 (reduce_sum.289: f32[], reduce_sum.290: f32[]) -> f32[] { @@ -312,39 +312,39 @@ StackFrames ROOT %reduce_sum.219 = f32[]{:T(128)} add(%reduce_sum.217, %reduce_sum.218), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.247 (param_0.1114: f32[14336,4,4096], param_1.1277: f32[], param_2.1108: f32[], param_3.797: f32[], param_4.498: f32[14336,4,4096], param_5.423: f32[], param_6.295: f32[4,14336,4096], param_7.194: pred[], param_8.112: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { - %param_0.1114 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %param_3.797 = f32[]{:T(128)S(6)} parameter(3) - %mul.1550.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.797), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.194 = pred[]{:T(512)S(6)} parameter(7) - %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.194), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.295 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) - %bitcast.423.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.295), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.423 = f32[]{:T(128)} parameter(5) - %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.423), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.423.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.423.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.248 (param_0.1100: f32[14336,4,4096], param_1.1263: f32[], param_2.1099: f32[], param_3.793: f32[], param_4.493: f32[14336,4,4096], param_5.415: f32[], param_6.283: f32[4,14336,4096], param_7.182: pred[], param_8.108: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { + %param_0.1100 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) + %param_3.793 = f32[]{:T(128)S(6)} parameter(3) + %mul.1616.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.182 = pred[]{:T(512)S(6)} parameter(7) + %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.182), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.283 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) + %bitcast.417.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.415 = f32[]{:T(128)} parameter(5) + %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.417.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.417.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.955.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.586.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.955.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1556.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.112 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) + %mul.1622.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.108 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) %constant.959.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1557.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1555.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.112, %mul.1557.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1556.clone.1, %mul.1555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1108 = f32[]{:T(128)S(6)} parameter(2) - %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1108), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1623.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1621.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.108, %mul.1623.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1622.clone.1, %mul.1621.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1099 = f32[]{:T(128)S(6)} parameter(2) + %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1099), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.68.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %select_n.273.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.958.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1554.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1552.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.498 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) + %mul.1620.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1618.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1620.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.493 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) %constant.957.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1553.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1551.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.498, %mul.1553.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1552.clone.1, %mul.1551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1277 = f32[]{:T(128)S(6)} parameter(1) - %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1619.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1617.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.493, %mul.1619.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1618.clone.1, %mul.1617.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1263 = f32[]{:T(128)S(6)} parameter(1) + %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1263), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.783.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.819.clone.1, %div.784.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.66.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} sqrt(%div.783.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.956.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -352,13 +352,13 @@ StackFrames %add.817.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%sqrt.66.clone.1, %add.818.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.264.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%div.785.clone.1, %add.817.clone.1), metadata={op_name="multiply.34"} %div.782.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.820.clone.1, %multiply.264.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1549.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1114, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1548.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1550.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1114, %mul.1548.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.171 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%add.815.clone.1, %add.815.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1615.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1100, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1615.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1614.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1616.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1100, %mul.1614.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.120 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%add.815.clone.1, %add.815.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1001 = f32[]{:T(128)} constant(0) - %reduce.127 = f32[]{:T(128)} reduce(%square.171, %constant.1001), dimensions={0,1,2}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.127 = f32[]{:T(128)} reduce(%square.120, %constant.1001), dimensions={0,1,2}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.130.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1001), dimensions={0,1,2}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.137 = (f32[]{:T(128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.127, %add.815.clone.1, %add.819.clone.1, %add.820.clone.1, %reduce.130.clone.1) } @@ -375,39 +375,39 @@ StackFrames ROOT %reduce_sum.213 = f32[]{:T(128)} add(%reduce_sum.211, %reduce_sum.212), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.248 (param_0.1115: f32[4096,4,14336], param_1.1278: f32[], param_2.1109: f32[], param_3.798: f32[], param_4.499: f32[4096,4,14336], param_5.424: f32[], param_6.296: f32[4,4096,14336], param_7.195: pred[], param_8.113: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1115 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.798 = f32[]{:T(128)S(6)} parameter(3) - %mul.1560.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.798), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.195 = pred[]{:T(512)S(6)} parameter(7) - %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.195), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.296 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.425.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.296), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.424 = f32[]{:T(128)} parameter(5) - %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.424), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.425.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.425.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.249 (param_0.1101: f32[4096,4,14336], param_1.1264: f32[], param_2.1100: f32[], param_3.794: f32[], param_4.494: f32[4096,4,14336], param_5.416: f32[], param_6.284: f32[4,4096,14336], param_7.183: pred[], param_8.109: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { + %param_0.1101 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.794 = f32[]{:T(128)S(6)} parameter(3) + %mul.1626.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.183 = pred[]{:T(512)S(6)} parameter(7) + %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.183), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.419.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.416 = f32[]{:T(128)} parameter(5) + %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.419.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.419.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.961.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.592.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.961.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1564.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.113 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) + %mul.1630.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.109 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) %constant.965.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.591.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.965.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1563.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.113, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1564.clone.1, %mul.1563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1109 = f32[]{:T(128)S(6)} parameter(2) - %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1109), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1629.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.109, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1630.clone.1, %mul.1629.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) + %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.69.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %select_n.277.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.964.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.590.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.964.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1562.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.499 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) + %mul.1628.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.494 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) %constant.963.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.589.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.963.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1561.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.499, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1562.clone.1, %mul.1561.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1278 = f32[]{:T(128)S(6)} parameter(1) - %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1278), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1627.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.494, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1628.clone.1, %mul.1627.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1264 = f32[]{:T(128)S(6)} parameter(1) + %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1264), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.791.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.824.clone.1, %div.792.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.67.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.791.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.962.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -415,13 +415,13 @@ StackFrames %add.823.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.67.clone.1, %broadcast.587.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.265.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.793.clone.1, %add.823.clone.1), metadata={op_name="multiply.33"} %div.790.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.825.clone.1, %multiply.265.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1559.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1115, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1558.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1560.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1115, %mul.1558.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.172 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.821.clone.1, %add.821.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1625.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1101, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1625.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1624.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1626.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1101, %mul.1624.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.121 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.821.clone.1, %add.821.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1002 = f32[]{:T(128)} constant(0) - %reduce.128 = f32[]{:T(128)} reduce(%square.172, %constant.1002), dimensions={0,1,2}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.128 = f32[]{:T(128)} reduce(%square.121, %constant.1002), dimensions={0,1,2}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.131.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1002), dimensions={0,1,2}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.138 = (f32[]{:T(128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.128, %add.821.clone.1, %add.824.clone.1, %add.825.clone.1, %reduce.131.clone.1) } @@ -438,39 +438,39 @@ StackFrames ROOT %reduce_sum.210 = f32[]{:T(128)} add(%reduce_sum.205, %reduce_sum.206), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.249 (param_0.1116: f32[4096,4,14336], param_1.1279: f32[], param_2.1110: f32[], param_3.799: f32[], param_4.500: f32[4096,4,14336], param_5.425: f32[], param_6.297: f32[4,4096,14336], param_7.196: pred[], param_8.114: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1116 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.799 = f32[]{:T(128)S(6)} parameter(3) - %mul.1567.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.799), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.196 = pred[]{:T(512)S(6)} parameter(7) - %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.196), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.297 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.427.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.297), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.425 = f32[]{:T(128)} parameter(5) - %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.425), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.427.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.427.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.250 (param_0.1102: f32[4096,4,14336], param_1.1265: f32[], param_2.1101: f32[], param_3.795: f32[], param_4.495: f32[4096,4,14336], param_5.417: f32[], param_6.285: f32[4,4096,14336], param_7.184: pred[], param_8.110: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { + %param_0.1102 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.795 = f32[]{:T(128)S(6)} parameter(3) + %mul.1633.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.184 = pred[]{:T(512)S(6)} parameter(7) + %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.184), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.285 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.421.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.417 = f32[]{:T(128)} parameter(5) + %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.421.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.421.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.967.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.598.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.967.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1571.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.114 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) + %mul.1637.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.110 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) %constant.971.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.597.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.971.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1570.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.114, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1571.clone.1, %mul.1570.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1110 = f32[]{:T(128)S(6)} parameter(2) - %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1110), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1636.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.110, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1637.clone.1, %mul.1636.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) + %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.70.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %select_n.281.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.970.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.596.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.970.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1569.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.500 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) + %mul.1635.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.495 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) %constant.969.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.595.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.969.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1568.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.500, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1569.clone.1, %mul.1568.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1279 = f32[]{:T(128)S(6)} parameter(1) - %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1279), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1634.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.495, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1635.clone.1, %mul.1634.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1265 = f32[]{:T(128)S(6)} parameter(1) + %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1265), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.799.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.829.clone.1, %div.800.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.68.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.799.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.968.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -478,13 +478,13 @@ StackFrames %add.828.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.68.clone.1, %broadcast.593.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.266.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.801.clone.1, %add.828.clone.1), metadata={op_name="multiply.32"} %div.798.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.830.clone.1, %multiply.266.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1566.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1116, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1565.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1567.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1116, %mul.1565.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.173 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.826.clone.1, %add.826.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1632.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1102, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1632.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1631.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1633.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1102, %mul.1631.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.122 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.826.clone.1, %add.826.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1003 = f32[]{:T(128)} constant(0) - %reduce.129 = f32[]{:T(128)} reduce(%square.173, %constant.1003), dimensions={0,1,2}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.129 = f32[]{:T(128)} reduce(%square.122, %constant.1003), dimensions={0,1,2}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.132.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1003), dimensions={0,1,2}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.139 = (f32[]{:T(128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.129, %add.826.clone.1, %add.829.clone.1, %add.830.clone.1, %reduce.132.clone.1) } @@ -495,12 +495,12 @@ StackFrames ROOT %reduce_sum.183 = f32[]{:T(128)} add(%reduce_sum.178, %reduce_sum.182), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.267 (param_0.1120: f32[4,4096,32,128]) -> f32[] { - %param_0.1120 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.329 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1120), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.176 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.329, %bitcast.329), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.268 (param_0.1106: f32[4,4096,32,128]) -> f32[] { + %param_0.1106 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.323 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1408 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.323, %bitcast.323), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1007 = f32[]{:T(128)} constant(0) - ROOT %reduce.133 = f32[]{:T(128)} reduce(%square.176, %constant.1007), dimensions={0,1,2,3}, to_apply=%region_30.35, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.133 = f32[]{:T(128)} reduce(%mul.1408, %constant.1007), dimensions={0,1,2,3}, to_apply=%region_30.35, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_29.34 (reduce_sum.175: f32[], reduce_sum.176: f32[]) -> f32[] { @@ -509,18 +509,18 @@ StackFrames ROOT %reduce_sum.177 = f32[]{:T(128)} add(%reduce_sum.175, %reduce_sum.176), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.269 (param_0.1121: f32[4,32,128,4096]) -> f32[] { - %param_0.1121 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.333 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1121), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.179 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.333, %bitcast.333), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.270 (param_0.1107: f32[4,32,128,4096]) -> f32[] { + %param_0.1107 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.327 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1411 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.327, %bitcast.327), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1008 = f32[]{:T(128)} constant(0) - ROOT %reduce.134 = f32[]{:T(128)} reduce(%square.179, %constant.1008), dimensions={0,1,2,3}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.134 = f32[]{:T(128)} reduce(%mul.1411, %constant.1008), dimensions={0,1,2,3}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.270 (param_0.748: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { - %param_0.748 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.748), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.334 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.271 (param_0.735: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { + %param_0.735 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) + %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.735), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} + ROOT %bitcast.328 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_57.62 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { @@ -535,39 +535,39 @@ StackFrames ROOT %reduce_sum.246 = f32[]{:T(128)} add(%reduce_sum.241, %reduce_sum.245), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.271 (param_0.1109: f32[4096,4,32,128], param_1.1272: f32[], param_2.1103: f32[], param_3.792: f32[], param_4.493: f32[4096,4,32,128], param_5.418: f32[], param_6.290: f32[4,4096,32,128], param_7.189: pred[], param_8.107: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { - %param_0.1109 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.792 = f32[]{:T(128)S(6)} parameter(3) - %mul.1509.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.189 = pred[]{:T(512)S(6)} parameter(7) - %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.189), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.290 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.413.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.290), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.418 = f32[]{:T(128)} parameter(5) - %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.413.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.413.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.272 (param_0.1095: f32[4096,4,32,128], param_1.1258: f32[], param_2.1094: f32[], param_3.788: f32[], param_4.488: f32[4096,4,32,128], param_5.410: f32[], param_6.278: f32[4,4096,32,128], param_7.177: pred[], param_8.103: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { + %param_0.1095 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.788 = f32[]{:T(128)S(6)} parameter(3) + %mul.1575.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.788), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.177 = pred[]{:T(512)S(6)} parameter(7) + %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.177), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.278 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.407.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.278), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.410 = f32[]{:T(128)} parameter(5) + %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.410), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.407.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.407.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.925.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.564.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.925.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %mul.1515.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.107 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1581.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.103 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) %constant.929.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1516.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1514.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.107, %mul.1516.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1515.clone.1, %mul.1514.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1103 = f32[]{:T(128)S(6)} parameter(2) - %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1103), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1582.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1580.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.103, %mul.1582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1581.clone.1, %mul.1580.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1094 = f32[]{:T(128)S(6)} parameter(2) + %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1094), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.63.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %select_n.253.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.928.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1513.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1511.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1513.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.493 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1579.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1577.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1579.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.488 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) %constant.927.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1512.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1510.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.493, %mul.1512.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1511.clone.1, %mul.1510.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) - %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1578.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1576.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.488, %mul.1578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1577.clone.1, %mul.1576.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1258 = f32[]{:T(128)S(6)} parameter(1) + %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1258), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.743.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.792.clone.1, %div.744.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.61.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} sqrt(%div.743.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.926.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -575,13 +575,13 @@ StackFrames %add.790.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%sqrt.61.clone.1, %add.791.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.259.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%div.745.clone.1, %add.790.clone.1), metadata={op_name="multiply.39"} %div.742.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.793.clone.1, %multiply.259.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1508.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1109, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1508.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1507.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1509.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1109, %mul.1507.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.180 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%add.788.clone.1, %add.788.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1574.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1095, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1574.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1573.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1575.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1095, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.123 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%add.788.clone.1, %add.788.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.996 = f32[]{:T(128)} constant(0) - %reduce.135 = f32[]{:T(128)} reduce(%square.180, %constant.996), dimensions={0,1,2,3}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.135 = f32[]{:T(128)} reduce(%square.123, %constant.996), dimensions={0,1,2,3}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.139.clone.1 = f32[]{:T(128)} reduce(%integer_pow.63.clone.1, %constant.996), dimensions={0,1,2,3}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.140 = (f32[]{:T(128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.135, %add.788.clone.1, %add.792.clone.1, %add.793.clone.1, %reduce.139.clone.1) } @@ -598,39 +598,39 @@ StackFrames ROOT %reduce_sum.240 = f32[]{:T(128)} add(%reduce_sum.238, %reduce_sum.239), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.272 (param_0.1110: f32[32,4,128,4096], param_1.1273: f32[], param_2.1104: f32[], param_3.793: f32[], param_4.494: f32[32,4,128,4096], param_5.419: f32[], param_6.291: f32[4,32,128,4096], param_7.190: pred[], param_8.108: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { - %param_0.1110 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %param_3.793 = f32[]{:T(128)S(6)} parameter(3) - %mul.1519.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.190 = pred[]{:T(512)S(6)} parameter(7) - %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.190), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.291 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.415.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.291), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.419 = f32[]{:T(128)} parameter(5) - %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.419), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.415.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.415.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.273 (param_0.1096: f32[32,4,128,4096], param_1.1259: f32[], param_2.1095: f32[], param_3.789: f32[], param_4.489: f32[32,4,128,4096], param_5.411: f32[], param_6.279: f32[4,32,128,4096], param_7.178: pred[], param_8.104: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { + %param_0.1096 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) + %param_3.789 = f32[]{:T(128)S(6)} parameter(3) + %mul.1585.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.178 = pred[]{:T(512)S(6)} parameter(7) + %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.178), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.279 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.409.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.411 = f32[]{:T(128)} parameter(5) + %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.411), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.409.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.409.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.931.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.566.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.931.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %mul.1525.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.108 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) + %mul.1591.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.104 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) %constant.935.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1526.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1524.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.108, %mul.1526.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1525.clone.1, %mul.1524.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1104 = f32[]{:T(128)S(6)} parameter(2) - %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1104), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1592.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1590.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.104, %mul.1592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1591.clone.1, %mul.1590.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1095 = f32[]{:T(128)S(6)} parameter(2) + %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1095), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.64.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %select_n.257.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.934.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1523.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1521.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1523.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.494 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) + %mul.1589.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1587.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.489 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) %constant.933.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1522.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1520.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.494, %mul.1522.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1521.clone.1, %mul.1520.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1273 = f32[]{:T(128)S(6)} parameter(1) - %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1273), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1588.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1586.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.489, %mul.1588.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1587.clone.1, %mul.1586.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1259 = f32[]{:T(128)S(6)} parameter(1) + %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1259), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.751.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.798.clone.1, %div.752.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.62.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} sqrt(%div.751.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.932.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -638,13 +638,13 @@ StackFrames %add.796.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%sqrt.62.clone.1, %add.797.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.260.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%div.753.clone.1, %add.796.clone.1), metadata={op_name="multiply.38"} %div.750.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.799.clone.1, %multiply.260.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1518.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1110, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1518.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1517.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1519.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1110, %mul.1517.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.181 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%add.794.clone.1, %add.794.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1584.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1096, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1584.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1583.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1585.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1096, %mul.1583.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.124 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%add.794.clone.1, %add.794.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.997 = f32[]{:T(128)} constant(0) - %reduce.136 = f32[]{:T(128)} reduce(%square.181, %constant.997), dimensions={0,1,2,3}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.136 = f32[]{:T(128)} reduce(%square.124, %constant.997), dimensions={0,1,2,3}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.140.clone.1 = f32[]{:T(128)} reduce(%integer_pow.64.clone.1, %constant.997), dimensions={0,1,2,3}, to_apply=%region_42.47, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.141 = (f32[]{:T(128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.136, %add.794.clone.1, %add.798.clone.1, %add.799.clone.1, %reduce.140.clone.1) } @@ -655,23 +655,23 @@ StackFrames ROOT %reduce_sum.267 = f32[]{:T(128)} add(%reduce_sum.262, %reduce_sum.266), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.279 (param_0.1129: bf16[4,128,128256], param_1.1288: f32[4,128], param_2.1115: s32[4,128], param_3.803: bf16[4,128]) -> f32[4,128] { - %param_2.1115 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1115), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.280 (param_0.1115: bf16[4,128,128256], param_1.1274: f32[4,128], param_2.1106: s32[4,128], param_3.799: bf16[4,128]) -> f32[4,128] { + %param_2.1106 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1106), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.25 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.24 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1129 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.950 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1129), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.803 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.803), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.950, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_0.1115 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.938 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1115), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_3.799 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) + %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.799), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.938, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1274 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.60 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %constant.1017 = f32[]{:T(128)} constant(0) %broadcast.511 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%constant.1017), dimensions={}, metadata={op_name="broadcast.83"} - %mul.1373 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1373, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1424 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1424, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_7.10 (reduce_sum.93: f32[], reduce_sum.94: f32[]) -> f32[] { @@ -680,12 +680,12 @@ StackFrames ROOT %reduce_sum.95 = f32[]{:T(128)} add(%reduce_sum.93, %reduce_sum.94), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.284 (param_0.1130: bf16[4,128,128256], param_1.1289: bf16[4,128]) -> f32[4,128] { - %param_0.1130 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.956 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1130), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1289 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1289), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.956, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.285 (param_0.1116: bf16[4,128,128256], param_1.1275: bf16[4,128]) -> f32[4,128] { + %param_0.1116 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.944 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1116), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1275 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1275), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.944, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.54 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} %constant.1018 = f32[]{:T(128)} constant(0) ROOT %reduce.138 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1018), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} @@ -703,23 +703,23 @@ StackFrames ROOT %reduce_sum.171 = f32[]{:T(128)} add(%reduce_sum.169, %reduce_sum.170), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.290 (param_0.1122: f32[4,4096,8,128], param_1.1282: f32[4,4096,8,128]) -> (f32[], f32[]) { - %param_0.1122 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.350 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1122), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.184 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.350, %bitcast.350), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.291 (param_0.1108: f32[4,4096,8,128], param_1.1268: f32[4,4096,8,128]) -> (f32[], f32[]) { + %param_0.1108 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.344 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1439 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.344, %bitcast.344), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1009 = f32[]{:T(128)} constant(0) - %reduce.141 = f32[]{:T(128)} reduce(%square.184, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1282 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(1) - %bitcast.354.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.187.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.354.clone.1, %bitcast.354.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.142.clone.1 = f32[]{:T(128)} reduce(%square.187.clone.1, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.141 = f32[]{:T(128)} reduce(%mul.1439, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1268 = f32[4,4096,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) + %bitcast.348.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1442.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.348.clone.1, %bitcast.348.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.142.clone.1 = f32[]{:T(128)} reduce(%mul.1442.clone.1, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.156 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.141, %reduce.142.clone.1) } -%fused_computation.293 (param_0.807: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.807 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.807), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.355 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.294 (param_0.794: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { + %param_0.794 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.794), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} + ROOT %bitcast.349 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_58.63 (reduce_sum.324: f32[], reduce_sum.325: f32[]) -> f32[] { @@ -734,39 +734,39 @@ StackFrames ROOT %reduce_sum.252 = f32[]{:T(128)} add(%reduce_sum.247, %reduce_sum.248), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.294 (param_0.1108: f32[4096,4,8,128], param_1.1271: f32[], param_2.1102: f32[], param_3.791: f32[], param_4.492: f32[4096,4,8,128], param_5.417: f32[], param_6.289: f32[4,4096,8,128], param_7.188: pred[], param_8.106: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1108 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.791 = f32[]{:T(128)S(6)} parameter(3) - %mul.1502.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.188 = pred[]{:T(512)S(6)} parameter(7) - %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.188), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.289 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.289), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.417 = f32[]{:T(128)} parameter(5) - %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.411.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.295 (param_0.1094: f32[4096,4,8,128], param_1.1257: f32[], param_2.1093: f32[], param_3.787: f32[], param_4.487: f32[4096,4,8,128], param_5.409: f32[], param_6.277: f32[4,4096,8,128], param_7.176: pred[], param_8.102: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { + %param_0.1094 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %param_3.787 = f32[]{:T(128)S(6)} parameter(3) + %mul.1568.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.787), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.176 = pred[]{:T(512)S(6)} parameter(7) + %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.176), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.277 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.405.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.277), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.409 = f32[]{:T(128)} parameter(5) + %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.405.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.405.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.919.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.562.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.919.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1506.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.106 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1572.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.923.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.561.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.923.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1505.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.106, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1506.clone.1, %mul.1505.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) - %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1571.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.102, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1572.clone.1, %mul.1571.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1093 = f32[]{:T(128)S(6)} parameter(2) + %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1093), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.62.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %select_n.249.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.922.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.560.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.922.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1504.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.492 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1570.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.487 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.921.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.559.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.921.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1503.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.492, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1504.clone.1, %mul.1503.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1271 = f32[]{:T(128)S(6)} parameter(1) - %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1271), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1569.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.487, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1570.clone.1, %mul.1569.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1257 = f32[]{:T(128)S(6)} parameter(1) + %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1257), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.735.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.786.clone.1, %div.736.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.60.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.735.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.920.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -774,15 +774,15 @@ StackFrames %add.785.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.60.clone.1, %broadcast.557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.258.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.737.clone.1, %add.785.clone.1), metadata={op_name="multiply.40"} %div.734.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.787.clone.1, %multiply.258.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1501.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1108, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1501.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1500.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1502.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1108, %mul.1500.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.188 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.783.clone.1, %add.783.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1567.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1094, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1566.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1568.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1094, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.125 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.783.clone.1, %add.783.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.995 = f32[]{:T(128)} constant(0) - %reduce.143 = f32[]{:T(128)} reduce(%square.188, %constant.995), dimensions={0,1,2,3}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.143 = f32[]{:T(128)} reduce(%square.125, %constant.995), dimensions={0,1,2,3}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.145.clone.1 = f32[]{:T(128)} reduce(%integer_pow.62.clone.1, %constant.995), dimensions={0,1,2,3}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) + ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) } %region_55.60 (reduce_sum.304: f32[], reduce_sum.308: f32[]) -> f32[] { @@ -797,39 +797,39 @@ StackFrames ROOT %reduce_sum.234 = f32[]{:T(128)} add(%reduce_sum.232, %reduce_sum.233), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.295 (param_0.1111: f32[4096,4,8,128], param_1.1274: f32[], param_2.1105: f32[], param_3.794: f32[], param_4.495: f32[4096,4,8,128], param_5.420: f32[], param_6.292: f32[4,4096,8,128], param_7.191: pred[], param_8.109: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1111 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.794 = f32[]{:T(128)S(6)} parameter(3) - %mul.1529.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.191 = pred[]{:T(512)S(6)} parameter(7) - %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.191), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.292 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.417.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.292), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.420 = f32[]{:T(128)} parameter(5) - %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.420), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.417.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.417.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.296 (param_0.1097: f32[4096,4,8,128], param_1.1260: f32[], param_2.1096: f32[], param_3.790: f32[], param_4.490: f32[4096,4,8,128], param_5.412: f32[], param_6.280: f32[4,4096,8,128], param_7.179: pred[], param_8.105: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { + %param_0.1097 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.790 = f32[]{:T(128)S(6)} parameter(3) + %mul.1595.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.179 = pred[]{:T(512)S(6)} parameter(7) + %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.179), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.280 = f32[4,4096,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.412 = f32[]{:T(128)} parameter(5) + %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.411.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.937.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.572.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.937.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1533.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.109 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1599.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.105 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.941.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.571.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.941.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1532.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.109, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1533.clone.1, %mul.1532.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1105 = f32[]{:T(128)S(6)} parameter(2) - %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1105), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1598.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.105, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1599.clone.1, %mul.1598.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1096 = f32[]{:T(128)S(6)} parameter(2) + %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1096), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.65.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %select_n.261.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.940.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.570.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.940.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1531.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.495 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1597.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.490 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.939.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.569.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.939.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1530.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.495, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1531.clone.1, %mul.1530.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1274 = f32[]{:T(128)S(6)} parameter(1) - %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1596.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.490, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1597.clone.1, %mul.1596.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1260 = f32[]{:T(128)S(6)} parameter(1) + %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1260), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.759.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.803.clone.1, %div.760.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.63.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.759.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.938.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -837,33 +837,33 @@ StackFrames %add.802.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.761.clone.1, %add.802.clone.1), metadata={op_name="multiply.37"} %div.758.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.804.clone.1, %multiply.261.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1528.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1111, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1528.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1527.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1529.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1111, %mul.1527.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.189 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.800.clone.1, %add.800.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1594.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1097, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1594.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1593.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1595.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1097, %mul.1593.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.126 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.800.clone.1, %add.800.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.998 = f32[]{:T(128)} constant(0) - %reduce.144 = f32[]{:T(128)} reduce(%square.189, %constant.998), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.144 = f32[]{:T(128)} reduce(%square.126, %constant.998), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.146.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.998), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.143 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.144, %add.800.clone.1, %add.803.clone.1, %add.804.clone.1, %reduce.146.clone.1) } -%fused_computation.311 (param_0.872: bf16[4,128,4096], param_1.941: f32[4,128], param_2.726: f32[4,128], param_3.452: bf16[4,128,4096], param_4.271: bf16[4096]) -> bf16[4,128,4096] { - %param_3.452 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.271 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.375 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.271), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.365 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.452, %dot_general.375), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.973 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.726 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1423 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.726), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1415 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.973, %mul.1423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.872 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.984 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.872), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.941 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1422 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.941), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1421 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.984, %mul.1422), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1415, %mul.1421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} - ROOT %convert_element_type.971 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.312 (param_0.859: bf16[4,128,4096], param_1.928: f32[4,128], param_2.717: f32[4,128], param_3.448: bf16[4,128,4096], param_4.266: bf16[4096]) -> bf16[4,128,4096] { + %param_3.448 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.266 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %dot_general.371 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.266), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.361 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.448, %dot_general.371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.961 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_2.717 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %mul.1480 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.717), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1472 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.961, %mul.1480), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_0.859 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.972 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.859), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.928 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1479 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.928), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1478 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.972, %mul.1479), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1472, %mul.1478), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} + ROOT %convert_element_type.959 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} } %region_5.8 (reduce_sum.87: f32[], reduce_sum.88: f32[]) -> f32[] { @@ -872,12 +872,12 @@ StackFrames ROOT %reduce_sum.92 = f32[]{:T(128)} add(%reduce_sum.87, %reduce_sum.88), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.312 (param_0.1131: bf16[4,128,4096]) -> f32[4,128] { - %param_0.1131 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.975 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1131), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.192 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.975, %convert_element_type.975), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} +%fused_computation.313 (param_0.1117: bf16[4,128,4096]) -> f32[4,128] { + %param_0.1117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.963 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1117), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.129 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.963, %convert_element_type.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} %constant.1019 = f32[]{:T(128)} constant(0) - ROOT %reduce.147 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.192, %constant.1019), dimensions={2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + ROOT %reduce.147 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.129, %constant.1019), dimensions={2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_10.13 (reduce_sum.102: f32[], reduce_sum.106: f32[]) -> f32[] { @@ -886,17 +886,17 @@ StackFrames ROOT %reduce_sum.107 = f32[]{:T(128)} add(%reduce_sum.102, %reduce_sum.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.314 (param_0.1126: bf16[4,128,4096], param_1.1285: bf16[4,128,4096], param_2.1113: bf16[4096]) -> f32[4,128] { - %param_0.1126 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1126), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1285 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.374 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1113), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.364 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1285, %dot_general.374), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.981 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1419 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %convert_element_type.981), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} +%fused_computation.315 (param_0.1112: bf16[4,128,4096], param_1.1271: bf16[4,128,4096], param_2.1104: bf16[4096]) -> f32[4,128] { + %param_0.1112 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.970 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1112), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1104 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.370 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1104), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.360 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1271, %dot_general.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.969 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.1476 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.970, %convert_element_type.969), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.1013 = f32[]{:T(128)} constant(0) - ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1419, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1476, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_8.11 (dot_general.182: bf16[], dot_general.183: bf16[]) -> bf16[] { @@ -905,86 +905,86 @@ StackFrames ROOT %add.168 = bf16[]{:T(256)} add(%dot_general.182, %dot_general.183), metadata={op_name="add.54"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.235.clone.clone (param_0.1095: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1095 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1033 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1095), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.449 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +%fused_computation.236.clone.clone (param_0.1081: f32[4096,128256]) -> bf16[4096,128256,1] { + %param_0.1081 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %convert_element_type.1021 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1081), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + ROOT %bitcast.443 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1021), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} } -%fused_computation.280.clone.1.clone.clone (param_0.1096: bf16[4,128,128256], param_1.1261: s32[4,128], param_2.1081: f32[4,128], param_3.782: f32[4,128], param_4.484: bf16[4,128], param_5.409: f32[4,128]) -> bf16[4,128,128256] { - %param_5.409 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1603 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.782 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1602 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.782), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1096 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1036 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1096), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.484 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.484), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1036, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.281.clone.1.clone.clone (param_0.1082: bf16[4,128,128256], param_1.1247: s32[4,128], param_2.1072: f32[4,128], param_3.778: f32[4,128], param_4.479: bf16[4,128], param_5.401: f32[4,128]) -> bf16[4,128,128256] { + %param_5.401 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.1669 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.401), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.778 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.1668 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.778), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1082 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1024 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.479 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.479), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1024, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.60 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1601 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1602, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1081), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1601, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1261 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1261), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.1667 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1668, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1072 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1072), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1667, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1247 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1247), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.42 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.41 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1035 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1035), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1600 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1603, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1034 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1600), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.315 (param_0.1094: f32[4,128], param_1.1260: bf16[4,128,4096], param_2.1082: f32[4096,128256], param_3.783: bf16[4,128,128256], param_4.485: s32[4,128], param_5.410: f32[4,128], param_6.284: f32[4,128], param_7.183: bf16[4,128], param_8.102: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { - %param_3.783 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.485 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.410 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.284 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.183 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.102 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.783, %param_4.485, %param_5.410, %param_6.284, %param_7.183, /*index=5*/%param_8.102), kind=kLoop, calls=%fused_computation.280.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1082 = f32[4096,128256]{1,0:T(8,128)} parameter(2) - %fusion.219.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1082), kind=kLoop, calls=%fused_computation.235.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.86.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.219.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %param_1.1260 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.994 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1260), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1094 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1434 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1094), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1433 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.994, %mul.1434), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.993 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1433), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.86.clone.1, %convert_element_type.993), metadata={op_name="multiply.206"} + %convert_element_type.1023 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1023), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1666 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1669, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1022 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1666), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.316 (param_0.1080: f32[4,128], param_1.1246: bf16[4,128,4096], param_2.1073: f32[4096,128256], param_3.779: bf16[4,128,128256], param_4.480: s32[4,128], param_5.402: f32[4,128], param_6.272: f32[4,128], param_7.171: bf16[4,128], param_8.98: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { + %param_3.779 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) + %param_4.480 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %param_5.402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.272 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.171 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) + %param_8.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) + %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.779, %param_4.480, %param_5.402, %param_6.272, %param_7.171, /*index=5*/%param_8.98), kind=kLoop, calls=%fused_computation.281.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_2.1073 = f32[4096,128256]{1,0:T(8,128)} parameter(2) + %fusion.209.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1073), kind=kLoop, calls=%fused_computation.236.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %convolution.80.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.209.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %param_1.1246 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1246), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1080 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1491 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1080), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1490 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %mul.1491), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.981 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1490), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.80.clone.1, %convert_element_type.981), metadata={op_name="multiply.206"} %constant.874 = bf16[]{:T(256)} constant(0) %reduce.149 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.252, %constant.874), dimensions={0,1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.86.clone.1) + ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.80.clone.1) } -%fused_computation.323 (param_0.904: f32[64], param_1.974: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.974), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.904 = f32[64]{0:T(128)S(1)} parameter(0) - %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.904), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} +%fused_computation.324 (param_0.891: f32[64], param_1.961: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.961 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.961), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %param_0.891 = f32[64]{0:T(128)S(1)} parameter(0) + %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.891), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %div.618 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.621, %div.619), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.618), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} - %convert_element_type.1002 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.990 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %cos.41.clone.1 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.618), metadata={op_name="jit(train_step)/layers/cos" stack_frame_id=0} - %convert_element_type.1001.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1002, %convert_element_type.1001.clone.1) + %convert_element_type.989.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.990, %convert_element_type.989.clone.1) } -%fused_computation.324 (param_0.901: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.901 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.325 (param_0.888: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.888 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.866 = bf16[]{:T(256)} constant(-inf) - %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.888, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.888, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.34 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.38, %pad.37), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } -%fused_computation.325 (param_0.903: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.903 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.326 (param_0.890: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.890 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.865 = bf16[]{:T(256)} constant(-inf) - %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.890, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.890, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.35 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.40, %pad.39), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } @@ -1000,16 +1000,16 @@ StackFrames ROOT %reduce_sum.162 = f32[]{:T(128)} add(%reduce_sum.157, %reduce_sum.161), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.329 (param_0.1123: f32[4,4096], param_1.1283: f32[4,4096]) -> (f32[], f32[]) { - %param_0.1123 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.371 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.195 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.330 (param_0.1109: f32[4,4096], param_1.1269: f32[4,4096]) -> (f32[], f32[]) { + %param_0.1109 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.365 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1500 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.365, %bitcast.365), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1010 = f32[]{:T(128)} constant(0) - %reduce.150 = f32[]{:T(128)} reduce(%square.195, %constant.1010), dimensions={0,1}, to_apply=%region_27.32, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1283 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) - %bitcast.375.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.198.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.151.clone.1 = f32[]{:T(128)} reduce(%square.198.clone.1, %constant.1010), dimensions={0,1}, to_apply=%region_26.31, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.150 = f32[]{:T(128)} reduce(%mul.1500, %constant.1010), dimensions={0,1}, to_apply=%region_27.32, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1269 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) + %bitcast.369.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1503.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.369.clone.1, %bitcast.369.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.151.clone.1 = f32[]{:T(128)} reduce(%mul.1503.clone.1, %constant.1010), dimensions={0,1}, to_apply=%region_26.31, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.150, %reduce.151.clone.1) } @@ -1025,39 +1025,39 @@ StackFrames ROOT %reduce_sum.231 = f32[]{:T(128)} add(%reduce_sum.226, %reduce_sum.227), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.332 (param_0.1112: f32[4096,4], param_1.1275: f32[], param_2.1106: f32[], param_3.795: f32[], param_4.496: f32[4096,4], param_5.421: f32[], param_6.293: f32[4,4096], param_7.192: pred[], param_8.110: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1112 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.795 = f32[]{:T(128)S(6)} parameter(3) - %mul.1536.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.192 = pred[]{:T(512)S(6)} parameter(7) - %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.192), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.293 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.419.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.293), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.421 = f32[]{:T(128)} parameter(5) - %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.421), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.419.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.419.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.333 (param_0.1098: f32[4096,4], param_1.1261: f32[], param_2.1097: f32[], param_3.791: f32[], param_4.491: f32[4096,4], param_5.413: f32[], param_6.281: f32[4,4096], param_7.180: pred[], param_8.106: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1098 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.791 = f32[]{:T(128)S(6)} parameter(3) + %mul.1602.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.180 = pred[]{:T(512)S(6)} parameter(7) + %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.180), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.281 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.413.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.281), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.413 = f32[]{:T(128)} parameter(5) + %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.413), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.413.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.413.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.943.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.578.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.943.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1540.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.110 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1606.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.106 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) %constant.947.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.577.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.947.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1539.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.110, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1540.clone.1, %mul.1539.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1106 = f32[]{:T(128)S(6)} parameter(2) - %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1106), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1605.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.106, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1606.clone.1, %mul.1605.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1097 = f32[]{:T(128)S(6)} parameter(2) + %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1097), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.66.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %select_n.265.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.946.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.576.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.946.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1538.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.496 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1604.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.491 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) %constant.945.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.575.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.945.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1537.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.496, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1538.clone.1, %mul.1537.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1275 = f32[]{:T(128)S(6)} parameter(1) - %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1275), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1603.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.491, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1604.clone.1, %mul.1603.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1261 = f32[]{:T(128)S(6)} parameter(1) + %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1261), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.767.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.808.clone.1, %div.768.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.64.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.767.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.944.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1065,13 +1065,13 @@ StackFrames %add.807.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.262.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.769.clone.1, %add.807.clone.1), metadata={op_name="multiply.36"} %div.766.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.809.clone.1, %multiply.262.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1535.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1112, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1535.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1534.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1536.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1112, %mul.1534.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.199 = f32[4096,4]{0,1:T(4,128)} multiply(%add.805.clone.1, %add.805.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1601.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1098, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1601.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1600.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1602.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1098, %mul.1600.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.130 = f32[4096,4]{0,1:T(4,128)} multiply(%add.805.clone.1, %add.805.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.999 = f32[]{:T(128)} constant(0) - %reduce.152 = f32[]{:T(128)} reduce(%square.199, %constant.999), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.152 = f32[]{:T(128)} reduce(%square.130, %constant.999), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.154.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.999), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.144 = (f32[]{:T(128)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.152, %add.805.clone.1, %add.808.clone.1, %add.809.clone.1, %reduce.154.clone.1) } @@ -1088,39 +1088,39 @@ StackFrames ROOT %reduce_sum.225 = f32[]{:T(128)} add(%reduce_sum.220, %reduce_sum.224), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.333 (param_0.1113: f32[4096,4], param_1.1276: f32[], param_2.1107: f32[], param_3.796: f32[], param_4.497: f32[4096,4], param_5.422: f32[], param_6.294: f32[4,4096], param_7.193: pred[], param_8.111: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1113 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.796 = f32[]{:T(128)S(6)} parameter(3) - %mul.1543.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.193 = pred[]{:T(512)S(6)} parameter(7) - %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.193), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.294 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.421.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.294), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.422 = f32[]{:T(128)} parameter(5) - %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.422), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.421.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.421.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.334 (param_0.1099: f32[4096,4], param_1.1262: f32[], param_2.1098: f32[], param_3.792: f32[], param_4.492: f32[4096,4], param_5.414: f32[], param_6.282: f32[4,4096], param_7.181: pred[], param_8.107: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1099 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.792 = f32[]{:T(128)S(6)} parameter(3) + %mul.1609.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.181 = pred[]{:T(512)S(6)} parameter(7) + %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.181), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.282 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.415.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.414 = f32[]{:T(128)} parameter(5) + %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.414), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.415.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.415.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.949.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.584.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.949.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1547.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.111 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1613.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.107 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) %constant.953.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.583.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.953.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1546.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.111, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1547.clone.1, %mul.1546.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1107 = f32[]{:T(128)S(6)} parameter(2) - %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1107), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1612.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.107, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1613.clone.1, %mul.1612.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1098 = f32[]{:T(128)S(6)} parameter(2) + %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1098), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.67.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %select_n.269.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.952.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.582.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.952.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1545.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.497 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1611.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.492 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) %constant.951.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.581.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.951.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1544.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.497, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1545.clone.1, %mul.1544.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1276 = f32[]{:T(128)S(6)} parameter(1) - %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1276), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1610.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.492, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1611.clone.1, %mul.1610.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1262 = f32[]{:T(128)S(6)} parameter(1) + %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1262), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.775.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.813.clone.1, %div.776.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.65.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.775.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.950.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1128,13 +1128,13 @@ StackFrames %add.812.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.65.clone.1, %broadcast.579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.263.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.777.clone.1, %add.812.clone.1), metadata={op_name="multiply.35"} %div.774.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.814.clone.1, %multiply.263.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1542.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1113, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1542.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1541.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1543.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1113, %mul.1541.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.200 = f32[4096,4]{0,1:T(4,128)} multiply(%add.810.clone.1, %add.810.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1608.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1099, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1608.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1607.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1609.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1099, %mul.1607.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.131 = f32[4096,4]{0,1:T(4,128)} multiply(%add.810.clone.1, %add.810.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1000 = f32[]{:T(128)} constant(0) - %reduce.153 = f32[]{:T(128)} reduce(%square.200, %constant.1000), dimensions={0,1}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.153 = f32[]{:T(128)} reduce(%square.131, %constant.1000), dimensions={0,1}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.155.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1000), dimensions={0,1}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.145 = (f32[]{:T(128)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.153, %add.810.clone.1, %add.813.clone.1, %add.814.clone.1, %reduce.155.clone.1) } @@ -1145,12 +1145,12 @@ StackFrames ROOT %reduce_sum.101 = f32[]{:T(128)} add(%reduce_sum.99, %reduce_sum.100), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.344 (param_0.1127: bf16[4096]) -> f32[] { - %param_0.1127 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.1006 = f32[4096]{0:T(1024)} convert(%param_0.1127), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.203 = f32[4096]{0:T(1024)} multiply(%convert_element_type.1006, %convert_element_type.1006), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.345 (param_0.1113: bf16[4096]) -> f32[] { + %param_0.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.994 = f32[4096]{0:T(1024)} convert(%param_0.1113), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.1520 = f32[4096]{0:T(1024)} multiply(%convert_element_type.994, %convert_element_type.994), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %constant.1014 = f32[]{:T(128)} constant(0) - ROOT %reduce.156 = f32[]{:T(128)} reduce(%square.203, %constant.1014), dimensions={0}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.156 = f32[]{:T(128)} reduce(%mul.1520, %constant.1014), dimensions={0}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_49.54 (reduce_sum.274: f32[], reduce_sum.275: f32[]) -> f32[] { @@ -1165,39 +1165,39 @@ StackFrames ROOT %reduce_sum.204 = f32[]{:T(128)} add(%reduce_sum.199, %reduce_sum.203), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.345 (param_0.1117: f32[4096], param_1.1280: f32[], param_2.1111: f32[], param_3.800: f32[], param_4.501: f32[4096], param_5.426: f32[], param_6.298: bf16[4096], param_7.197: pred[], param_8.115: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { - %param_0.1117 = f32[4096]{0:T(1024)S(1)} parameter(0) - %param_3.800 = f32[]{:T(128)S(6)} parameter(3) - %mul.1574.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.800), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.197 = pred[]{:T(512)S(6)} parameter(7) - %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.298 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1021.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.298), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.426 = f32[]{:T(128)} parameter(5) - %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.426), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1021.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1021.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.346 (param_0.1103: f32[4096], param_1.1266: f32[], param_2.1102: f32[], param_3.796: f32[], param_4.496: f32[4096], param_5.418: f32[], param_6.286: bf16[4096], param_7.185: pred[], param_8.111: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { + %param_0.1103 = f32[4096]{0:T(1024)S(1)} parameter(0) + %param_3.796 = f32[]{:T(128)S(6)} parameter(3) + %mul.1640.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.185 = pred[]{:T(512)S(6)} parameter(7) + %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.185), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.286 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1009.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.286), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_5.418 = f32[]{:T(128)} parameter(5) + %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1009.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1009.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.973.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.600.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.973.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.1580.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.115 = f32[4096]{0:T(1024)S(1)} parameter(8) + %mul.1646.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.111 = f32[4096]{0:T(1024)S(1)} parameter(8) %constant.977.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1581.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1579.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.115, %mul.1581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1580.clone.1, %mul.1579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1111 = f32[]{:T(128)S(6)} parameter(2) - %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1111), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1647.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1645.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.111, %mul.1647.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1646.clone.1, %mul.1645.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) + %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.71.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %select_n.285.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.976.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1578.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1576.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.501 = f32[4096]{0:T(1024)S(1)} parameter(4) + %mul.1644.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1642.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1644.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.496 = f32[4096]{0:T(1024)S(1)} parameter(4) %constant.975.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1577.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1575.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.501, %mul.1577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1576.clone.1, %mul.1575.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1280 = f32[]{:T(128)S(6)} parameter(1) - %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1280), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1643.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1641.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.496, %mul.1643.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1642.clone.1, %mul.1641.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1266 = f32[]{:T(128)S(6)} parameter(1) + %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1266), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.807.clone.1 = f32[4096]{0:T(1024)} divide(%add.835.clone.1, %div.808.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.69.clone.1 = f32[4096]{0:T(1024)} sqrt(%div.807.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.974.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1205,38 +1205,38 @@ StackFrames %add.833.clone.1 = f32[4096]{0:T(1024)} add(%sqrt.69.clone.1, %add.834.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.267.clone.1 = f32[4096]{0:T(1024)} multiply(%div.809.clone.1, %add.833.clone.1), metadata={op_name="multiply.31"} %div.806.clone.1 = f32[4096]{0:T(1024)} divide(%add.836.clone.1, %multiply.267.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1573.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1117, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1572.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1574.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1117, %mul.1572.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.204 = f32[4096]{0:T(1024)} multiply(%add.831.clone.1, %add.831.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1639.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1103, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1639.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1638.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1640.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1103, %mul.1638.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.132 = f32[4096]{0:T(1024)} multiply(%add.831.clone.1, %add.831.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1004 = f32[]{:T(128)} constant(0) - %reduce.157 = f32[]{:T(128)} reduce(%square.204, %constant.1004), dimensions={0}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.157 = f32[]{:T(128)} reduce(%square.132, %constant.1004), dimensions={0}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.158.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1004), dimensions={0}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.148 = (f32[]{:T(128)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.157, %add.831.clone.1, %add.835.clone.1, %add.836.clone.1, %reduce.158.clone.1) } -%fused_computation.351 (param_0.964: s32[512]) -> s32[1024] { +%fused_computation.352 (param_0.951: s32[512]) -> s32[1024] { %constant.801 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.539 = s32[1024]{0:T(1024)} broadcast(%constant.801), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.964 = s32[512]{0:T(512)S(1)} parameter(0) + %param_0.951 = s32[512]{0:T(512)S(1)} parameter(0) %constant.802 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.964, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.951, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %constant.800 = s32[] constant(128255), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.538 = s32[1024]{0:T(1024)} broadcast(%constant.800), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.539, %pad.41, %broadcast.538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } -%fused_computation.352 (param_0.963: s32[4,128]) -> s32[512] { - %param_0.963 = s32[4,128]{1,0:T(4,128)} parameter(0) +%fused_computation.353 (param_0.950: s32[4,128]) -> s32[512] { + %param_0.950 = s32[4,128]{1,0:T(4,128)} parameter(0) %constant.888 = s32[]{:T(128)} constant(0) %broadcast.546 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.888), dimensions={}, metadata={op_name="broadcast.81"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.963, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} + %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.950, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} %constant.875 = s32[]{:T(128)} constant(128256) %add.760 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.875), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.963, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} - ROOT %bitcast.376 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.950, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} + %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} + ROOT %bitcast.370 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} } %region_61.66 (reduce_sum.345: f32[], reduce_sum.346: f32[]) -> f32[] { @@ -1251,52 +1251,52 @@ StackFrames ROOT %reduce_sum.273 = f32[]{:T(128)} add(%reduce_sum.268, %reduce_sum.269), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.353 (param_0.1128: bf16[4,128], param_1.1287: f32[4,128], param_2.1114: f32[4,128], param_3.802: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { - %param_3.802 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) +%fused_computation.354 (param_0.1114: bf16[4,128], param_1.1273: f32[4,128], param_2.1105: f32[4,128], param_3.798: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { + %param_3.798 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) %constant.979.clone.1 = s32[]{:T(128)} constant(0) %broadcast.601.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.979.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.802, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1287), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1128 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1128), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.798, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %param_1.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1273), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_0.1114 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1114), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} %add.762 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %square.207 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %add.762), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %square.135 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %add.762), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} %constant.1016 = f32[]{:T(128)} constant(0) %broadcast.543 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1016), dimensions={}, metadata={op_name="broadcast.32"} - %mul.1473 = f32[4,128]{1,0:T(4,128)} multiply(%square.207, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %mul.1465 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1473, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.159 = f32[]{:T(128)} reduce(%mul.1465, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1114), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} - %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1473), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %mul.1466.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1466.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %mul.1471.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1539 = f32[4,128]{1,0:T(4,128)} multiply(%square.135, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1531 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1539, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.159 = f32[]{:T(128)} reduce(%mul.1531, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_2.1105 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1105), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1539), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1532.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1532.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1537.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %constant.891.clone.1 = f32[]{:T(128)} constant(1) %add.757.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.891.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1471.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1537.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} ROOT %tuple.149 = (f32[]{:T(128)}, f32[]{:T(128)}, pred[4,128]{1,0:T(4,128)(4,1)S(1)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%reduce.159, %reduce.160.clone.1, %ne.6.clone.1, %add.750.clone.1) } -%fused_computation.356 (param_0.987: f32[4,128], param_1.1101: f32[4,128]) -> f32[4,128] { - %param_0.987 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1101 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.357 (param_0.974: f32[4,128], param_1.1088: f32[4,128]) -> f32[4,128] { + %param_0.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1088 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.869 = f32[]{:T(128)} constant(0.000244140625) %broadcast.549 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.869), dimensions={}, metadata={op_name="broadcast.264"} - %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1101, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1088, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.867 = f32[]{:T(128)} constant(1e-05) %add.770 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.867), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.769 = f32[4,128]{1,0:T(4,128)} add(%div.656, %add.770), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %rsqrt.90 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} %div.649 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.90, %add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.864 = f32[]{:T(128)} constant(-0.5) - %mul.1477 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1470 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1477), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1469 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.987, %mul.1470), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1543 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1536 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1543), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1535 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.974, %mul.1536), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.863 = f32[]{:T(128)} constant(0.00048828125) - %mul.1476 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - ROOT %mul.1468 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1469, %mul.1476), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1542 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + ROOT %mul.1534 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1535, %mul.1542), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} } %region_0.1 (reduce_sum.67: s32[], reduce_sum.71: s32[]) -> s32[] { @@ -1305,64 +1305,64 @@ StackFrames ROOT %reduce_sum.72 = s32[]{:T(128)} add(%reduce_sum.67, %reduce_sum.71), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.360 (param_0.1004: pred[4,128]) -> s32[] { - %param_0.1004 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1013 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1004), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} +%fused_computation.361 (param_0.991: pred[4,128]) -> s32[] { + %param_0.991 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %convert_element_type.1001 = s32[4,128]{1,0:T(4,128)} convert(%param_0.991), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} %constant.889 = s32[]{:T(128)} constant(0) - ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1013, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1001, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.361 (param_0.989: f32[4,128]) -> f32[4,128] { - %param_0.989 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.362 (param_0.976: f32[4,128]) -> f32[4,128] { + %param_0.976 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.870 = f32[]{:T(128)} constant(0.000244140625) %broadcast.541 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.870), dimensions={}, metadata={op_name="broadcast.264"} - %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.989, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.976, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.868 = f32[]{:T(128)} constant(1e-05) %add.759 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.868), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.756 = f32[4,128]{1,0:T(4,128)} add(%div.654, %add.759), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} ROOT %rsqrt.88 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.756), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} } -%fused_computation.362 (param_0.990: pred[4,128], param_1.1286: f32[]) -> f32[4,128] { - %param_0.990 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1286 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1286), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} +%fused_computation.363 (param_0.977: pred[4,128], param_1.1272: f32[]) -> f32[4,128] { + %param_0.977 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) + %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} %constant.1015 = f32[]{:T(128)} constant(0) %broadcast.545 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1015), dimensions={}, metadata={op_name="broadcast.32"} - ROOT %mul.1478 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.990, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %mul.1544 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.977, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} } -%fused_computation.364 () -> f32[64] { +%fused_computation.365 () -> f32[64] { %constant.873 = f32[]{:T(128)} constant(500000) %broadcast.552 = f32[64]{0:T(128)} broadcast(%constant.873), dimensions={}, metadata={op_name="broadcast.255"} %iota.46 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/layers/iota" stack_frame_id=0} %constant.872 = s32[]{:T(128)} constant(2) %broadcast.551 = s32[64]{0:T(128)} broadcast(%constant.872), dimensions={}, metadata={op_name="broadcast.256"} - %mul.1479 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} - %convert_element_type.1014 = f32[64]{0:T(128)} convert(%mul.1479), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %mul.1545 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} + %convert_element_type.1002 = f32[64]{0:T(128)} convert(%mul.1545), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %constant.871 = f32[]{:T(128)} constant(0.0078125) %broadcast.550 = f32[64]{0:T(128)} broadcast(%constant.871), dimensions={}, metadata={op_name="broadcast.257"} - %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1014, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1002, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.552, %div.657), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} } -%fused_computation.365 (param_0.1002: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.1002 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1015 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1002), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %bitcast.377 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1015), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.377, %convert_element_type.1015) +%fused_computation.366 (param_0.989: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { + %param_0.989 = s32[4,128]{1,0:T(4,128)} parameter(0) + %convert_element_type.1003 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.989), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %bitcast.371 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.371, %convert_element_type.1003) } -%fused_computation.369 (param_0.1103: f32[4096,4]) -> bf16[4,4096] { - %param_0.1103 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.451 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1103), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.451) +%fused_computation.369 (param_0.1089: f32[4096,4]) -> bf16[4,4096] { + %param_0.1089 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.445 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1089), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.445) } -%fused_computation.370 (param_0.1104: f32[4096,4]) -> bf16[4,4096] { - %param_0.1104 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.452 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1104), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.452) +%fused_computation.370 (param_0.1090: f32[4096,4]) -> bf16[4,4096] { + %param_0.1090 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.446 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.446) } %region_6.9 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { @@ -1371,41 +1371,41 @@ StackFrames ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.237.clone.clone (param_0.1090: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1090 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1026 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.447 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1026), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} -} - -%fused_computation.317.clone.clone (param_0.1091: f32[4,128], param_1.1257: bf16[4,128,4096], param_2.1077: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1077 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1077), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1257 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1028 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1257), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1091 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1595 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1091), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1594 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1028, %mul.1595), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1027 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1594), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1027), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.371 (param_0.1105: f32[4096,128256], param_1.1268: f32[4,128], param_2.1099: bf16[4,128,4096], param_3.788: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { - %param_1.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1099 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.788 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.240.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1268, %param_2.1099, %param_3.788), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1105 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %fusion.221.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1105), kind=kLoop, calls=%fused_computation.237.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.87.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.240.clone.1, %fusion.221.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} +%fused_computation.238.clone.clone (param_0.1076: f32[4096,128256]) -> bf16[4096,128256,1] { + %param_0.1076 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %convert_element_type.1014 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1076), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + ROOT %bitcast.441 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1014), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +} + +%fused_computation.318.clone.clone (param_0.1077: f32[4,128], param_1.1243: bf16[4,128,4096], param_2.1068: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1068 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.379 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1068), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1243 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1016 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1243), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1077 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1661 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1077), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1660 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1016, %mul.1661), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1015 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1660), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.378 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.379, %convert_element_type.1015), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.371 (param_0.1091: f32[4096,128256], param_1.1254: f32[4,128], param_2.1090: bf16[4,128,4096], param_3.784: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { + %param_1.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1090 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.784 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.230.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1254, %param_2.1090, %param_3.784), kind=kLoop, calls=%fused_computation.318.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1091 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %fusion.211.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1091), kind=kLoop, calls=%fused_computation.238.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %convolution.81.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.230.clone.1, %fusion.211.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} %constant.992 = bf16[]{:T(256)} constant(-inf) - %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.87.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} - ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.87.clone.1) + %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.81.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.81.clone.1) } -%fused_computation.372 (param_0.1102: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.1102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %bitcast.450 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1102), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.450) +%fused_computation.372 (param_0.1088: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { + %param_0.1088 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %bitcast.444 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1088), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.444) } %convert_element_type.525.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { @@ -1414,13 +1414,13 @@ StackFrames ROOT %add.624 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.121.clone.clone (param_0.1242: bf16[4,4096], param_1.1376: s32[]) -> bf16[4096] { - %param_0.1242 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.121.clone.clone (param_0.1229: bf16[4,4096], param_1.1363: s32[]) -> bf16[4096] { + %param_0.1229 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1363 = s32[]{:T(128)S(6)} parameter(1) %constant.1116 = s32[]{:T(128)} constant(0) - %dynamic_slice.316 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1242, %param_1.1376, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %dynamic_slice.310 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1229, %param_1.1363, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} %constant.1117 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.316, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.310, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %region_12.14 (reduce_sum.108: f32[], reduce_sum.109: f32[]) -> f32[] { @@ -1429,70 +1429,70 @@ StackFrames ROOT %reduce_sum.113 = f32[]{:T(128)} add(%reduce_sum.108, %reduce_sum.109), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.58.clone.clone (param_0.1243: bf16[4,4,128,4096], param_1.1377: s32[]) -> f32[4,128] { - %param_0.1243 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1377 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.58.clone.clone (param_0.1230: bf16[4,4,128,4096], param_1.1364: s32[]) -> f32[4,128] { + %param_0.1230 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1364 = s32[]{:T(128)S(6)} parameter(1) %constant.1118 = s32[]{:T(128)} constant(0) - %dynamic_slice.317 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1243, %param_1.1377, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.548 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1093 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.548), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.214 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1093, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %dynamic_slice.311 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1230, %param_1.1364, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.543 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.311), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1081 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.543), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.142 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1081, %convert_element_type.1081), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} %constant.1119 = f32[]{:T(128)} constant(0) - ROOT %reduce.175 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.214, %constant.1119), dimensions={2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + ROOT %reduce.175 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.142, %constant.1119), dimensions={2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} } -%fused_computation.143.clone.1.clone (param_0.1244: f32[4,128]) -> f32[4,128] { - %param_0.1244 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.143.clone.1.clone (param_0.1231: f32[4,128]) -> f32[4,128] { + %param_0.1231 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1121 = f32[]{:T(128)} constant(0.000244140625) %closed_call.81 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1121), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1244, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1231, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1120 = f32[]{:T(128)} constant(1e-05) %closed_call.80 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1120), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.858 = f32[4,128]{1,0:T(4,128)} add(%div.842, %closed_call.80), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.97 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.858), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.24.clone.1.clone.clone (param_0.1258: bf16[4,4096,32,128], param_1.1387: s32[]) -> bf16[4096,32,128,1] { - %param_0.1258 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.24.clone.1.clone.clone (param_0.1245: bf16[4,4096,32,128], param_1.1374: s32[]) -> bf16[4096,32,128,1] { + %param_0.1245 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1374 = s32[]{:T(128)S(6)} parameter(1) %constant.1134 = s32[]{:T(128)} constant(0) - %dynamic_slice.323 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1258, %param_1.1387, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.559 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.317 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1245, %param_1.1374, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.554 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.91.clone.clone (param_0.1259: f32[4,128], param_1.1388: bf16[4,4,128,4096], param_2.1176: s32[], param_3.847: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.847 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.847), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1388 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1176 = s32[]{:T(128)S(6)} parameter(2) +%fused_computation.91.clone.clone (param_0.1246: f32[4,128], param_1.1375: bf16[4,4,128,4096], param_2.1167: s32[], param_3.843: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.843 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.843), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1375 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1167 = s32[]{:T(128)S(6)} parameter(2) %constant.1135 = s32[]{:T(128)} constant(0) - %dynamic_slice.324 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1388, %param_2.1176, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.561 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1101 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1259 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1709 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1259), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1708 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1101, %mul.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1100 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1708), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1100), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.560 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.427), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.36.clone.clone (param_0.1260: bf16[4,4096,32,128], param_1.1389: s32[], param_2.1177: f32[4,128], param_3.848: bf16[4,4,128,4096], param_4.530: bf16[4096]) -> bf16[4,128,32,128] { - %param_2.1177 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.848 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1389 = s32[]{:T(128)S(6)} parameter(1) - %param_4.530 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.343 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1177, %param_3.848, %param_1.1389, %param_4.530), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1260 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.342 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1260, %param_1.1389), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.113 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.343, %fusion.342), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.70.clone.clone (param_0.1261: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { - %param_0.1261 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %dynamic_slice.318 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1375, %param_2.1167, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.556 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.318), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1089 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.556), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1246 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1775 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1246), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1774 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1089, %mul.1775), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1088 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1774), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1088), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.555 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.36.clone.clone (param_0.1247: bf16[4,4096,32,128], param_1.1376: s32[], param_2.1168: f32[4,128], param_3.844: bf16[4,4,128,4096], param_4.525: bf16[4096]) -> bf16[4,128,32,128] { + %param_2.1168 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.844 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) + %param_4.525 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.332 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1168, %param_3.844, %param_1.1376, %param_4.525), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1247 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.331 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1247, %param_1.1376), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.107 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.332, %fusion.331), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.70.clone.clone (param_0.1248: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { + %param_0.1248 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1248), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.129 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1248), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.187 = (bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) } @@ -1502,172 +1502,172 @@ StackFrames %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} %constant.1123 = s32[]{:T(128)} constant(2) %closed_call.83 = s32[64]{0:T(128)} broadcast(%constant.1123), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1699 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1094 = f32[64]{0:T(128)} convert(%mul.1699), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %mul.1765 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1082 = f32[64]{0:T(128)} convert(%mul.1765), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %constant.1122 = f32[]{:T(128)} constant(0.0078125) %closed_call.82 = f32[64]{0:T(128)} broadcast(%constant.1122), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1094, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1082, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.84, %div.843), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} } -%fused_computation.117.clone.clone (param_0.1245: f32[64], param_1.1378: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1378 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1378), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1245 = f32[64]{0:T(128)S(1)} parameter(0) - %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1245), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} +%fused_computation.117.clone.clone (param_0.1232: f32[64], param_1.1365: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1365 = f32[4,128]{1,0:T(4,128)} parameter(1) + %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1365), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %param_0.1232 = f32[64]{0:T(128)S(1)} parameter(0) + %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1232), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %div.844 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.846, %div.845), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1095 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1083 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} %convert_element_type.829.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1095, %convert_element_type.829.clone.3) + ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1083, %convert_element_type.829.clone.3) } -%fused_computation.120.clone.clone (param_0.1252: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1252 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.120.clone.clone (param_0.1239: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1239 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1130 = bf16[]{:T(256)} constant(-inf) - %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1239, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1239, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.61, %pad.60), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.554 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.119.clone.clone (param_0.1246: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1246 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.119.clone.clone (param_0.1233: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1233 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1125 = bf16[]{:T(256)} constant(-inf) - %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1233, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1233, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.44 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.59, %pad.58), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.544 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.73.clone.clone (param_0.1262: bf16[4,128,32,64], param_1.1390: bf16[4,128,32,64], param_2.1178: bf16[4,128,32,128], param_3.849: bf16[4,128,128], param_4.531: bf16[4,128,128]) -> bf16[4,32,128,128] { - %param_2.1178 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.531 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1713 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.531), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1711 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1178, %mul.1713), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1390 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.73.clone.clone (param_0.1249: bf16[4,128,32,64], param_1.1377: bf16[4,128,32,64], param_2.1169: bf16[4,128,32,128], param_3.845: bf16[4,128,128], param_4.526: bf16[4,128,128]) -> bf16[4,32,128,128] { + %param_2.1169 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.526 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1779 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.526), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1777 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1169, %mul.1779), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1377 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1136 = bf16[]{:T(256)} constant(-inf) - %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1390, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1262 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1262, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1377, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1249 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1249, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.47 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.65, %pad.64), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.849 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1712 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.849), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1710 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1712), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1711, %mul.1710), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.562 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.90.clone.clone (param_0.1254: f32[4,128], param_1.1384: bf16[4,4,128,4096], param_2.1173: s32[], param_3.844: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.844 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.844), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1384 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1173 = s32[]{:T(128)S(6)} parameter(2) + %param_3.845 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1778 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.845), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1776 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1778), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1777, %mul.1776), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.557 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.90.clone.clone (param_0.1241: f32[4,128], param_1.1371: bf16[4,4,128,4096], param_2.1164: s32[], param_3.840: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.840 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.422 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.840), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1371 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1164 = s32[]{:T(128)S(6)} parameter(2) %constant.1132 = s32[]{:T(128)} constant(0) - %dynamic_slice.322 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1384, %param_2.1173, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.557 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1099 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1703 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1254), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1702 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1099, %mul.1703), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1098 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1098), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.556 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.425), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.64.clone.1.clone.clone (param_0.1253: bf16[4,4096,8,128], param_1.1383: s32[]) -> bf16[4096,8,128,1] { - %param_0.1253 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1383 = s32[]{:T(128)S(6)} parameter(1) + %dynamic_slice.316 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1371, %param_2.1164, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.316), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1087 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1241 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1769 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1241), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1768 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1087, %mul.1769), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1086 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1768), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.421 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.422, %convert_element_type.1086), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.64.clone.1.clone.clone (param_0.1240: bf16[4,4096,8,128], param_1.1370: s32[]) -> bf16[4096,8,128,1] { + %param_0.1240 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1370 = s32[]{:T(128)S(6)} parameter(1) %constant.1131 = s32[]{:T(128)} constant(0) - %dynamic_slice.321 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1383, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.555 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.315 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1240, %param_1.1370, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.550 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.315), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.89.clone.clone (param_0.1255: bf16[4,4096,8,128], param_1.1385: s32[], param_2.1174: f32[4,128], param_3.845: bf16[4,4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,8,128] { - %param_2.1174 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.845 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) - %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.340 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1174, %param_3.845, %param_1.1385, %param_4.528), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1255 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.341 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1255, %param_1.1385), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.112 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.340, %fusion.341), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.89.clone.clone (param_0.1242: bf16[4,4096,8,128], param_1.1372: s32[], param_2.1165: f32[4,128], param_3.841: bf16[4,4,128,4096], param_4.523: bf16[4096]) -> bf16[4,128,8,128] { + %param_2.1165 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.841 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1372 = s32[]{:T(128)S(6)} parameter(1) + %param_4.523 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.329 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1165, %param_3.841, %param_1.1372, %param_4.523), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1242 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.330 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1242, %param_1.1372), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.329, %fusion.330), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } -%fused_computation.106.clone.clone (param_0.1256: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1256 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.106.clone.clone (param_0.1243: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { + %param_0.1243 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1243), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1243), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.186 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) } -%fused_computation.109.clone.clone (param_0.1257: bf16[4,128,8,64], param_1.1386: bf16[4,128,8,64], param_2.1175: bf16[4,128,8,128], param_3.846: bf16[4,128,128], param_4.529: bf16[4,128,128]) -> bf16[4,8,128,128] { - %param_2.1175 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.529 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1707 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.529), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1705 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1175, %mul.1707), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1386 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.109.clone.clone (param_0.1244: bf16[4,128,8,64], param_1.1373: bf16[4,128,8,64], param_2.1166: bf16[4,128,8,128], param_3.842: bf16[4,128,128], param_4.524: bf16[4,128,128]) -> bf16[4,8,128,128] { + %param_2.1166 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.524 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1773 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.524), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1771 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1166, %mul.1773), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1373 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1133 = bf16[]{:T(256)} constant(-inf) - %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1386, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1257 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1257, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1373, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1244 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1244, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.46 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.63, %pad.62), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.846 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1706 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.846), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1704 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1706), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1705, %mul.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.558 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_3.842 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1772 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1770 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1772), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1771, %mul.1770), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.135.clone.clone (param_0.1248: bf16[4,4096,8,128], param_1.1380: s32[]) -> bf16[1,4096,8,128] { - %param_0.1248 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.135.clone.clone (param_0.1235: bf16[4,4096,8,128], param_1.1367: s32[]) -> bf16[1,4096,8,128] { + %param_0.1235 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) + %param_1.1367 = s32[]{:T(128)S(6)} parameter(1) %constant.1128 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.319 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1248, %param_1.1380, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %dynamic_slice.313 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1235, %param_1.1367, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } -%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1249: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { - %param_0.1249 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.550 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1236: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { + %param_0.1236 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1236), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} + ROOT %bitcast.545 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.88.clone.clone.clone.clone (param_0.1250: f32[4,128], param_1.1381: bf16[4,4,128,4096], param_2.1171: s32[], param_3.842: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.842 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1381 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) +%fused_computation.88.clone.clone.clone.clone (param_0.1237: f32[4,128], param_1.1368: bf16[4,4,128,4096], param_2.1162: s32[], param_3.838: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.838 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.420 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.838), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1368 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1162 = s32[]{:T(128)S(6)} parameter(2) %constant.1129 = s32[]{:T(128)} constant(0) - %dynamic_slice.320 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1381, %param_2.1171, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1097 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1250 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1701 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1250), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1700 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1097, %mul.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1096 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1700), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1096), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.114.clone.clone (param_0.1251: bf16[1,4096,8,128], param_1.1382: f32[4,128], param_2.1172: bf16[4,4,128,4096], param_3.843: s32[], param_4.527: bf16[4096]) -> bf16[4,8,128,128] { - %param_1.1382 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1172 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.843 = s32[]{:T(128)S(6)} parameter(3) - %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.339 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1382, %param_2.1172, %param_3.843, %param_4.527), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1251 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.338 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1251), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.111 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.339, %fusion.338), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.366.clone.clone (param_0.1286: f32[4,32,128,128]) -> (f32[4,32,128,1], f32[4,32,128]) { - %param_0.1286 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1286), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.262.clone.3 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.192 = (f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}, f32[4,32,128]{2,1,0:T(8,128)S(1)}) tuple(%slice.11, %bitcast.262.clone.3) + %dynamic_slice.314 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1368, %param_2.1162, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.547 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.314), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1085 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.547), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1237 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1767 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1237), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1766 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1085, %mul.1767), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1084 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1766), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.419 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.420, %convert_element_type.1084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.546 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.419), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.114.clone.clone (param_0.1238: bf16[1,4096,8,128], param_1.1369: f32[4,128], param_2.1163: bf16[4,4,128,4096], param_3.839: s32[], param_4.522: bf16[4096]) -> bf16[4,8,128,128] { + %param_1.1369 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1163 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.839 = s32[]{:T(128)S(6)} parameter(3) + %param_4.522 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.328 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1369, %param_2.1163, %param_3.839, %param_4.522), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1238 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.327 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1238), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.105 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.328, %fusion.327), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + ROOT %bitcast.548 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.146.clone.clone (param_0.1273: f32[4,32,128,128]) -> (f32[4,32,128], f32[4,32,128,1]) { + %param_0.1273 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) + %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1273), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.570 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} + ROOT %tuple.192 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.570, %slice.11) } %region_13.16 (reduce_sum.120: f32[], reduce_sum.121: f32[]) -> f32[] { @@ -1676,36 +1676,36 @@ StackFrames ROOT %reduce_sum.122 = f32[]{:T(128)} add(%reduce_sum.120, %reduce_sum.121), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1263: bf16[4,32,128,4096], param_1.1391: s32[]) -> bf16[32,128,4096,1] { - %param_0.1263 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1250: bf16[4,32,128,4096], param_1.1378: s32[]) -> bf16[32,128,4096,1] { + %param_0.1250 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1378 = s32[]{:T(128)S(6)} parameter(1) %constant.1137 = s32[]{:T(128)} constant(0) - %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1263, %param_1.1391, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.563 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.319 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1250, %param_1.1378, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.558 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.319), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1264: bf16[4,32,128,128]) -> bf16[4,128,32,128] { - %param_0.1264 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.564 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1251: bf16[4,32,128,128]) -> bf16[4,128,32,128] { + %param_0.1251 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.559 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.61.clone.clone (param_0.1265: bf16[4,32,128,4096], param_1.1392: s32[], param_2.1179: bf16[4,32,128,128], param_3.850: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { - %param_3.850 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1392 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.61.clone.clone (param_0.1252: bf16[4,32,128,4096], param_1.1379: s32[], param_2.1170: bf16[4,32,128,128], param_3.846: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { + %param_3.846 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) %constant.365.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.208.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.850, %param_1.1392, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.208.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1179 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.83.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1179), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1265 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.82.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1265, %param_1.1392), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.62.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.83.clone.3, %fusion.82.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %dynamic_slice.210.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.846, %param_1.1379, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.210.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %param_2.1170 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.80.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1170), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_0.1252 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.79.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1252, %param_1.1379), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.60.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.80.clone.3, %fusion.79.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %add.635.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.207.clone.3, %bitcast.182.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1102 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.215 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1102, %convert_element_type.1102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %convert_element_type.1090 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.143 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1090, %convert_element_type.1090), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} %constant.1138 = f32[]{:T(128)} constant(0) - %reduce.177 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.215, %constant.1138), dimensions={2}, to_apply=%region_13.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %reduce.177 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.143, %constant.1138), dimensions={2}, to_apply=%region_13.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.188 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.177, %add.635.clone.3) } @@ -1715,140 +1715,140 @@ StackFrames ROOT %add.623 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.122.clone.clone (param_0.1247: bf16[4,4096], param_1.1379: s32[]) -> bf16[4096] { - %param_0.1247 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.122.clone.clone (param_0.1234: bf16[4,4096], param_1.1366: s32[]) -> bf16[4096] { + %param_0.1234 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1366 = s32[]{:T(128)S(6)} parameter(1) %constant.1126 = s32[]{:T(128)} constant(0) - %dynamic_slice.318 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1247, %param_1.1379, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %dynamic_slice.312 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1234, %param_1.1366, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} %constant.1127 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.318, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.312, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.12.clone.clone.clone (param_0.1266: bf16[4,14336,4096], param_1.1393: s32[]) -> bf16[14336,4096,1] { - %param_0.1266 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1393 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.12.clone.clone.clone (param_0.1253: bf16[4,14336,4096], param_1.1380: s32[]) -> bf16[14336,4096,1] { + %param_0.1253 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) %constant.1139 = s32[]{:T(128)} constant(0) - %dynamic_slice.326 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1266, %param_1.1393, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.566 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.326), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.320 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1380, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.561 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %bitcast_fusion.3.clone.clone (bitcast_input.12: bf16[4,128,4096]) -> bf16[4,128,4096] { %bitcast_input.12 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.565 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) + ROOT %bitcast.560 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) } -%fused_computation.13.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1394: bf16[4,14336,4096], param_2.1180: s32[]) -> bf16[14336,4,128] { - %param_1.1394 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) - %fusion.344 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1394, %param_2.1180), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %fusion.345 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1267), kind=kLoop, calls=%bitcast_fusion.3.clone.clone - ROOT %convolution.114 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.344, %fusion.345), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.13.clone.clone (param_0.1254: bf16[4,128,4096], param_1.1381: bf16[4,14336,4096], param_2.1171: s32[]) -> bf16[14336,4,128] { + %param_1.1381 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) + %fusion.333 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1381, %param_2.1171), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1254 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.334 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1254), kind=kLoop, calls=%bitcast_fusion.3.clone.clone + ROOT %convolution.108 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.333, %fusion.334), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} } -%fused_computation.144.clone.1.clone (param_0.1268: f32[4,128]) -> f32[4,128] { - %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.144.clone.1.clone (param_0.1255: f32[4,128]) -> f32[4,128] { + %param_0.1255 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1141 = f32[]{:T(128)} constant(0.000244140625) %closed_call.86 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1141), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1255, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1140 = f32[]{:T(128)} constant(1e-05) %closed_call.85 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1140), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.861 = f32[4,128]{1,0:T(4,128)} add(%div.847, %closed_call.85), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.98 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.11.clone.1.clone.clone (param_0.1272: bf16[4,4096,14336], param_1.1398: s32[]) -> bf16[4096,14336,1] { - %param_0.1272 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1398 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.11.clone.1.clone.clone (param_0.1259: bf16[4,4096,14336], param_1.1385: s32[]) -> bf16[4096,14336,1] { + %param_0.1259 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) %constant.1143 = s32[]{:T(128)} constant(0) - %dynamic_slice.328 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1272, %param_1.1398, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.568 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.322 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1259, %param_1.1385, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.563 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.96.clone.2.clone.clone (param_0.1273: f32[4,128], param_1.1399: bf16[4,128,4096], param_2.1183: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1183 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.432 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1183), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1399 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1106 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1717 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1273), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1716 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1106, %mul.1717), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1105 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1716), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.431 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.432, %convert_element_type.1105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.96.clone.2.clone.clone (param_0.1260: f32[4,128], param_1.1386: bf16[4,128,4096], param_2.1174: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1174), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1094 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1260 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1783 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1260), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1782 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1094, %mul.1783), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1093 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1782), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.23.clone.clone (param_0.1274: bf16[4,4096,14336], param_1.1400: s32[], param_2.1184: f32[4,128], param_3.852: bf16[4,128,4096], param_4.533: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1184 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.533 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.349 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1184, %param_3.852, %param_4.533), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1274 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1400 = s32[]{:T(128)S(6)} parameter(1) - %fusion.348 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1274, %param_1.1400), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.116 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.349, %fusion.348), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.23.clone.clone (param_0.1261: bf16[4,4096,14336], param_1.1387: s32[], param_2.1175: f32[4,128], param_3.848: bf16[4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,14336] { + %param_2.1175 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.848 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.338 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1175, %param_3.848, %param_4.528), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1261 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) + %fusion.337 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1261, %param_1.1387), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.110 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.338, %fusion.337), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } -%fused_computation.14.clone.1.clone.clone (param_0.1275: bf16[4,4096,14336], param_1.1401: s32[]) -> bf16[4096,14336,1] { - %param_0.1275 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1401 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.14.clone.1.clone.clone (param_0.1262: bf16[4,4096,14336], param_1.1388: s32[]) -> bf16[4096,14336,1] { + %param_0.1262 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1388 = s32[]{:T(128)S(6)} parameter(1) %constant.1144 = s32[]{:T(128)} constant(0) - %dynamic_slice.329 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1275, %param_1.1401, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.569 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.329), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.323 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1262, %param_1.1388, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.564 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.39.clone.1.clone.clone (param_0.1276: bf16[14336,4,128], param_1.1402: bf16[4,128,14336]) -> bf16[4,128,14336] { - %param_1.1402 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.39.clone.1.clone.clone (param_0.1263: bf16[14336,4,128], param_1.1389: bf16[4,128,14336]) -> bf16[4,128,14336] { + %param_1.1389 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1145 = bf16[]{:T(256)} constant(1) %jit_silu_.44 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1145), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1402), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.69 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.862 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.848 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.862), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1719 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1402, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %param_0.1276 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) - %bitcast.570 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - ROOT %mul.1718 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1719, %bitcast.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1785 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1389, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %param_0.1263 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.565 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + ROOT %mul.1784 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1785, %bitcast.565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} } -%fused_computation.21.clone.clone (param_0.1277: bf16[4,4096,14336], param_1.1403: s32[], param_2.1185: bf16[14336,4,128], param_3.853: bf16[4,128,14336]) -> bf16[4,128,4096] { - %param_2.1185 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %param_3.853 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1185, %param_3.853), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_0.1277 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1403 = s32[]{:T(128)S(6)} parameter(1) - %fusion.350 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1277, %param_1.1403), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.350), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.21.clone.clone (param_0.1264: bf16[4,4096,14336], param_1.1390: s32[], param_2.1176: bf16[14336,4,128], param_3.849: bf16[4,128,14336]) -> bf16[4,128,4096] { + %param_2.1176 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) + %param_3.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1176, %param_3.849), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %param_0.1264 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1390 = s32[]{:T(128)S(6)} parameter(1) + %fusion.339 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1264, %param_1.1390), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.111 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.339), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} } -%fused_computation.14.clone.clone.clone (param_0.1269: bf16[4,4096,14336], param_1.1395: s32[]) -> bf16[4096,14336,1] { - %param_0.1269 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1395 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.14.clone.clone.clone (param_0.1256: bf16[4,4096,14336], param_1.1382: s32[]) -> bf16[4096,14336,1] { + %param_0.1256 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1382 = s32[]{:T(128)S(6)} parameter(1) %constant.1142 = s32[]{:T(128)} constant(0) - %dynamic_slice.327 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1269, %param_1.1395, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.567 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.321 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1256, %param_1.1382, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.562 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.96.clone.1.clone.clone (param_0.1270: f32[4,128], param_1.1396: bf16[4,128,4096], param_2.1181: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1181 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1181), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1396 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1104 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1270 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1715 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1270), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1714 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1104, %mul.1715), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1103 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1714), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.430, %convert_element_type.1103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.96.clone.1.clone.clone (param_0.1257: f32[4,128], param_1.1383: bf16[4,128,4096], param_2.1172: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1172 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1172), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1092 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1257 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1781 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1257), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1780 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1092, %mul.1781), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1091 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1780), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1091), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.20.clone.clone (param_0.1271: bf16[4,4096,14336], param_1.1397: s32[], param_2.1182: f32[4,128], param_3.851: bf16[4,128,4096], param_4.532: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1182 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.532 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.347 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1182, %param_3.851, %param_4.532), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1271 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1397 = s32[]{:T(128)S(6)} parameter(1) - %fusion.346 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1271, %param_1.1397), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.115 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.347, %fusion.346), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.20.clone.clone (param_0.1258: bf16[4,4096,14336], param_1.1384: s32[], param_2.1173: f32[4,128], param_3.847: bf16[4,128,4096], param_4.527: bf16[4096]) -> bf16[4,128,14336] { + %param_2.1173 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.847 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.336 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1173, %param_3.847, %param_4.527), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1258 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1384 = s32[]{:T(128)S(6)} parameter(1) + %fusion.335 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1258, %param_1.1384), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.109 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.336, %fusion.335), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } %region_14.17 (reduce_sum.126: f32[], reduce_sum.127: f32[]) -> f32[] { @@ -1857,63 +1857,63 @@ StackFrames ROOT %reduce_sum.128 = f32[]{:T(128)} add(%reduce_sum.126, %reduce_sum.127), metadata={op_name="checkpoint/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1278: bf16[4,4096,14336], param_1.1404: s32[]) -> bf16[4096,14336,1] { - %param_0.1278 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1404 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1265: bf16[4,4096,14336], param_1.1391: s32[]) -> bf16[4096,14336,1] { + %param_0.1265 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) %constant.1146 = s32[]{:T(128)} constant(0) - %dynamic_slice.330 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1278, %param_1.1404, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.571 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.324 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1265, %param_1.1391, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.566 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1279: bf16[4,128,14336], param_1.1405: bf16[4,128,14336], param_2.1186: bf16[14336,4,128]) -> bf16[4,128,14336] { - %param_2.1186 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %bitcast.572 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1186), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %param_1.1405 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %mul.1724 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.572, %param_1.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} +%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1266: bf16[4,128,14336], param_1.1392: bf16[4,128,14336], param_2.1177: bf16[14336,4,128]) -> bf16[4,128,14336] { + %param_2.1177 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.567 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1177), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + %param_1.1392 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %mul.1790 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.567, %param_1.1392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1147 = bf16[]{:T(256)} constant(1) %jit_silu_.45 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1147), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %param_0.1279 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %param_0.1266 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1266), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.70 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.131), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.863 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.70, %jit_silu_.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.45, %add.863), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1723 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1724, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1722 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1279, %mul.1724), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1789 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1790, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1788 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1266, %mul.1790), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} %sub.98 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} subtract(%jit_silu_.45, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/sub" stack_frame_id=0} - %mul.1721 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1720 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1722, %mul.1721), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1723, %mul.1720), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} -} - -%fused_computation.63.clone.clone (param_0.1280: bf16[4,128,4096], param_1.1406: bf16[4096], param_2.1187: bf16[4,128,4096], param_3.854: bf16[4,4096,14336], param_4.534: s32[], param_5.435: bf16[4,128,14336], param_6.304: bf16[4,128,14336], param_7.200: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { - %param_0.1280 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1108 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_2.1187 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_5.435 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) - %param_6.304 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) - %param_7.200 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) - %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.435, %param_6.304, %param_7.200), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} - %param_3.854 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.534 = s32[]{:T(128)S(6)} parameter(4) - %fusion.79.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.854, %param_4.534), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.60.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.79.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1187, %convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1406 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) - %dot_general.434 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1406), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.433 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.434), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1107 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.433), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1725 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1108, %convert_element_type.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1787 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1786 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1788, %mul.1787), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1789, %mul.1786), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} +} + +%fused_computation.63.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1393: bf16[4096], param_2.1178: bf16[4,128,4096], param_3.850: bf16[4,4096,14336], param_4.529: s32[], param_5.425: bf16[4,128,14336], param_6.291: bf16[4,128,14336], param_7.188: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { + %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1096 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1267), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_2.1178 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_5.425 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) + %param_6.291 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) + %param_7.188 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) + %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.425, %param_6.291, %param_7.188), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} + %param_3.850 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) + %param_4.529 = s32[]{:T(128)S(6)} parameter(4) + %fusion.91.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.850, %param_4.529), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.64.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.91.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1178, %convolution.64.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %param_1.1393 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1393), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.430), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1095 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.429), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} + %mul.1791 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1096, %convert_element_type.1095), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1148 = f32[]{:T(128)} constant(0) - %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1725, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} + %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1791, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.189 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.178, %add_any.132.clone.3) } -%fused_computation.140.clone.clone (param_0.1281: f32[4,128], param_1.1407: f32[4,128]) -> f32[4,128] { - %param_0.1281 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1407 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.140.clone.clone (param_0.1268: f32[4,128], param_1.1394: f32[4,128]) -> f32[4,128] { + %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1394 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.1152 = f32[]{:T(128)} constant(0.000244140625) %closed_call.89 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1152), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1407, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1394, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1151 = f32[]{:T(128)} constant(1e-05) %closed_call.88 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1151), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.864 = f32[4,128]{1,0:T(4,128)} add(%div.851, %closed_call.88), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} @@ -1921,11 +1921,11 @@ StackFrames %div.850 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.99, %add.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1150 = f32[]{:T(128)} constant(-0.5) %closed_call.87 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1150), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1728 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1727 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1281, %mul.1728), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1794 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1793 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %mul.1794), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1149 = f32[]{:T(128)} constant(0.00048828125) - %mul.1729 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - ROOT %mul.1726 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1727, %mul.1729), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1795 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + ROOT %mul.1792 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1793, %mul.1795), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} } %region_20.24 (dot_general.187: bf16[], dot_general.188: bf16[]) -> bf16[] { @@ -1934,29 +1934,29 @@ StackFrames ROOT %add.173 = bf16[]{:T(256)} add(%dot_general.187, %dot_general.188), metadata={op_name="add.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.94.clone.clone (param_0.1282: bf16[4,128,4096], param_1.1408: f32[4,128], param_2.1188: bf16[4,128,4096], param_3.855: bf16[4,128,4096], param_4.535: f32[4,128], param_5.436: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { - %param_0.1282 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_2.1188 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %convert_element_type.1110 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1188), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_1.1408 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1731 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1408), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1730 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1109 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1730), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %convert_element_type.1109), metadata={op_name="multiply.204"} +%fused_computation.94.clone.clone (param_0.1269: bf16[4,128,4096], param_1.1395: f32[4,128], param_2.1179: bf16[4,128,4096], param_3.851: bf16[4,128,4096], param_4.530: f32[4,128], param_5.426: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { + %param_0.1269 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_2.1179 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1098 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1179), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_1.1395 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1797 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1395), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1796 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1098, %mul.1797), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1097 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1796), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1269, %convert_element_type.1097), metadata={op_name="multiply.204"} %constant.1153 = bf16[]{:T(256)} constant(0) %reduce.179 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.271, %constant.1153), dimensions={0,1}, to_apply=%region_20.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.855 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_5.436 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) - %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.436), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_5.426 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.426), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1269, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} %convert_element_type.753.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.263.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1142.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_4.535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %mul.1151.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.535), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %mul.1141.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1151.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %add_any.126.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1142.clone.3, %mul.1141.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %mul.1178.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1797), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %param_4.530 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %mul.1187.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.530), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1177.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1098, %mul.1187.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %add_any.126.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1178.clone.3, %mul.1177.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} %convert_element_type.751.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.126.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.855, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.851, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} ROOT %tuple.190 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.179, %add_any.124.clone.3) } @@ -1966,35 +1966,35 @@ StackFrames ROOT %add.169 = f32[]{:T(128)} add(%dot_general.184, %dot_general.185), metadata={op_name="add.31"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1283: bf16[4,32,128,4096], param_1.1409: s32[]) -> bf16[32,128,4096,1] { - %param_0.1283 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1409 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1270: bf16[4,32,128,4096], param_1.1396: s32[]) -> bf16[32,128,4096,1] { + %param_0.1270 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1396 = s32[]{:T(128)S(6)} parameter(1) %constant.1154 = s32[]{:T(128)} constant(0) - %dynamic_slice.331 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1283, %param_1.1409, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.573 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1270, %param_1.1396, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.568 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1284: bf16[4,128,4096]) -> bf16[4,128,4096,1] { - %param_0.1284 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.574 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} +%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1271: bf16[4,128,4096]) -> bf16[4,128,4096,1] { + %param_0.1271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.569 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1271), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} } -%fused_computation.66.clone.clone (param_0.1285: bf16[4,32,128,128], param_1.1410: bf16[4,32,128,4096], param_2.1189: s32[], param_3.856: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { - %param_0.1285 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} - %param_3.856 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %fusion.95.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.856), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1410 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1189 = s32[]{:T(128)S(6)} parameter(2) - %fusion.94.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1410, %param_2.1189), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.64.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.95.clone.3, %fusion.94.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.66.clone.clone (param_0.1272: bf16[4,32,128,128], param_1.1397: bf16[4,32,128,4096], param_2.1180: s32[], param_3.852: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { + %param_0.1272 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1272), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} + %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %fusion.84.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.852), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %param_1.1397 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) + %fusion.83.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1397, %param_2.1180), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.62.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.84.clone.3, %fusion.83.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} %constant.619.clone.3 = bf16[]{:T(256)} constant(0.25) %div.442.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.619.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} - %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.64.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} + %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.62.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} %bitcast.209.clone.3 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%div.441.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} %convert.123 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.209.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert.1" stack_frame_id=0} %multiply.272 = f32[4,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.124, %convert.123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/multiply" stack_frame_id=0} %constant.1155 = f32[]{:T(128)} constant(0) - %dot_general.435 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} - ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.435, %bitcast.209.clone.3) + %dot_general.431 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} + ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.431, %bitcast.209.clone.3) } diff --git a/tests/utils/reference_hlo_qwen3_1.7b.txt b/tests/utils/reference_hlo_qwen3_1.7b.txt index f1ede66966..4004648b77 100644 --- a/tests/utils/reference_hlo_qwen3_1.7b.txt +++ b/tests/utils/reference_hlo_qwen3_1.7b.txt @@ -32,7 +32,7 @@ StackFrames %param_1.5 = s32[512]{0:T(512)S(1)} parameter(1) %reshape.451 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.5), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} %transpose.466 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.451), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)} parameter(2) %reshape.452 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} reshape(%param_2.4), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} %transpose.467 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} transpose(%reshape.452), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} ROOT %scatter.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.466, %transpose.467), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_42.47.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} @@ -50,43 +50,43 @@ StackFrames ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.386, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.277 (param_0.1368: f32[151936,2048], param_1.1556: f32[], param_2.1314: f32[], param_3.918: f32[], param_4.556: f32[151936,2048], param_5.468: f32[], param_6.358: bf16[151936,2048], param_7.201: bf16[151936,2048,1], param_8.118: pred[], param_9.97: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { - %param_0.1368 = f32[151936,2048]{1,0:T(8,128)} parameter(0) +%fused_computation.277 (param_0.1367: f32[151936,2048], param_1.1549: f32[], param_2.1311: f32[], param_3.918: f32[], param_4.554: f32[151936,2048], param_5.467: f32[], param_6.356: bf16[151936,2048], param_7.196: bf16[151936,2048,1], param_8.113: pred[], param_9.94: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { + %param_0.1367 = f32[151936,2048]{1,0:T(8,128)} parameter(0) %param_3.918 = f32[]{:T(128)S(6)} parameter(3) - %mul.1926.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_3.918), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.118 = pred[]{:T(512)S(6)} parameter(8) - %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.118), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_7.201 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) - %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.201), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %mul.2002.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_3.918), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.113 = pred[]{:T(512)S(6)} parameter(8) + %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.113), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_7.196 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) + %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.196), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %convert_element_type.1409.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.464.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_6.358 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_6.356 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.356), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %add_any.197.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1409.clone.1, %convert_element_type.1408.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} - %param_5.468 = f32[]{:T(128)} parameter(5) - %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_5.467 = f32[]{:T(128)} parameter(5) + %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.467), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.859.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add_any.197.clone.1, %div.860.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.267.clone.1 = f32[151936,2048]{1,0:T(8,128)} select(%select_n.268.clone.1, %add_any.197.clone.1, %div.859.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1092.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.844.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1092.clone.1), dimensions={}, metadata={op_name="broadcast.74"} - %mul.1932.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_9.97 = f32[151936,2048]{1,0:T(8,128)} parameter(9) + %mul.2008.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_9.94 = f32[151936,2048]{1,0:T(8,128)} parameter(9) %constant.1096.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1933.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1096.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1931.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.97, %mul.1933.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.941.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1932.clone.1, %mul.1931.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) - %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2009.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1096.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2007.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.94, %mul.2009.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.941.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.2008.clone.1, %mul.2007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1311 = f32[]{:T(128)S(6)} parameter(2) + %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1311), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.65.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %select_n.267.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1095.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1930.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1095.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1928.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %mul.1930.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.556 = f32[151936,2048]{1,0:T(8,128)} parameter(4) + %mul.2006.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1095.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2004.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %mul.2006.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.554 = f32[151936,2048]{1,0:T(8,128)} parameter(4) %constant.1094.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1929.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1094.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1927.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.556, %mul.1929.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.940.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1928.clone.1, %mul.1927.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) - %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2005.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1094.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2003.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.554, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.940.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.2004.clone.1, %mul.2003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1549 = f32[]{:T(128)S(6)} parameter(1) + %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1549), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.854.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.940.clone.1, %div.855.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.62.clone.1 = f32[151936,2048]{1,0:T(8,128)} sqrt(%div.854.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1093.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -94,14 +94,14 @@ StackFrames %add.938.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%sqrt.62.clone.1, %add.939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.426.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%div.856.clone.1, %add.938.clone.1), metadata={op_name="multiply.61"} %div.853.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.941.clone.1, %multiply.426.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1925.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1368, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.937.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%div.853.clone.1, %mul.1925.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1924.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%mul.1926.clone.1, %add.937.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1368, %mul.1924.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.214 = f32[151936,2048]{1,0:T(8,128)} multiply(%add.936.clone.1, %add.936.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1200 = f32[]{:T(128)} constant(0) - %reduce.176 = f32[]{:T(128)} reduce(%square.214, %constant.1200), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2001.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1367, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.937.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%div.853.clone.1, %mul.2001.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2000.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%mul.2002.clone.1, %add.937.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1367, %mul.2000.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.175 = f32[151936,2048]{1,0:T(8,128)} multiply(%add.936.clone.1, %add.936.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1198 = f32[]{:T(128)} constant(0) + %reduce.176 = f32[]{:T(128)} reduce(%square.175, %constant.1198), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1198), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.144 = (f32[]{:T(128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.176, %add.936.clone.1, %add.940.clone.1, %add.941.clone.1, %reduce.178.clone.1) } @@ -111,64 +111,64 @@ StackFrames ROOT %reduce_sum.319 = f32[]{:T(128)} add(%reduce_sum.317, %reduce_sum.318), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.367.clone.clone (param_0.1355: f32[4,128], param_1.1549: bf16[4,128,2048], param_2.1290: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1290 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.480 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1290), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1549 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1451 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1355 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2083 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1355), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2082 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1451, %mul.2083), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.480, %convert_element_type.1450), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.367.clone.clone (param_0.1354: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1287: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1287 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1287), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1445 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1354 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2151 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1354), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.2150 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1445, %mul.2151), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1444 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2150), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.478 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.479, %convert_element_type.1444), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.289.clone.clone.clone (param_0.1356: bf16[4,128,151936], param_1.1550: s32[4,128], param_2.1291: f32[4,128], param_3.911: f32[4,128], param_4.546: bf16[4,128], param_5.446: f32[4,128]) -> bf16[4,128,151936] { - %param_5.446 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2087 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.446), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.289.clone.clone.clone (param_0.1355: bf16[4,128,151936], param_1.1543: s32[4,128], param_2.1288: f32[4,128], param_3.911: f32[4,128], param_4.544: bf16[4,128], param_5.445: f32[4,128]) -> bf16[4,128,151936] { + %param_5.445 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.2155 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.445), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2086 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1356 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1454 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1356), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.546 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1454, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %mul.2154 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1355 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1448 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1355), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.544 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.544), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1448, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.62 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2085 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2086, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1291 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1291), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2085, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1550 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.2153 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2154, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2153, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1543), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.48 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.47 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1453 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1453), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2084 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2087, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1452 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.281 (param_0.1381: bf16[151936,2048], param_1.1569: f32[4,128], param_2.1327: bf16[4,128,2048], param_3.931: bf16[2048], param_4.569: bf16[4,128,151936], param_5.481: s32[4,128], param_6.371: f32[4,128], param_7.214: f32[4,128], param_8.131: bf16[4,128], param_9.98: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { - %param_4.569 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) - %param_5.481 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.371 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.214 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) - %param_8.131 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) - %param_9.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.569, %param_5.481, %param_6.371, %param_7.214, %param_8.131, /*index=5*/%param_9.98), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1327 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1447 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1447), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2152 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2155, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1446 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2152), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.281 (param_0.1380: bf16[151936,2048], param_1.1562: f32[4,128], param_2.1324: bf16[4,128,2048], param_3.931: bf16[2048], param_4.567: bf16[4,128,151936], param_5.480: s32[4,128], param_6.369: f32[4,128], param_7.209: f32[4,128], param_8.126: bf16[4,128], param_9.95: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { + %param_4.567 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) + %param_5.480 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.369 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.209 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) + %param_8.126 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) + %param_9.95 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.567, %param_5.480, %param_6.369, %param_7.209, %param_8.126, /*index=5*/%param_9.95), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1562 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1324 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) %param_3.931 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1569, %param_2.1327, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.269.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %fusion.268.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1562, %param_2.1324, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.268.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %bitcast.333 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%convolution.86.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %convert_element_type.1323 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_0.1381 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_0.1380 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1380), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %add_any.184 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1323, %convert_element_type.1322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} - %square.215 = f32[151936,2048]{1,0:T(8,128)} multiply(%add_any.184, %add_any.184), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1213 = f32[]{:T(128)} constant(0) - %reduce.177 = f32[]{:T(128)} reduce(%square.215, %constant.1213), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1800 = f32[151936,2048]{1,0:T(8,128)} multiply(%add_any.184, %add_any.184), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1211 = f32[]{:T(128)} constant(0) + %reduce.177 = f32[]{:T(128)} reduce(%mul.1800, %constant.1211), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.166 = (f32[]{:T(128)}, bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.177, %convolution.86.clone.1) } @@ -178,23 +178,23 @@ StackFrames ROOT %reduce_sum.394 = f32[]{:T(128)} add(%reduce_sum.389, %reduce_sum.393), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.288 (param_0.1392: bf16[4,128,151936], param_1.1577: f32[4,128], param_2.1330: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { - %param_2.1330 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1330), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.288 (param_0.1391: bf16[4,128,151936], param_1.1570: f32[4,128], param_2.1327: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { + %param_2.1327 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1327), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.25 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.24 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_0.1391 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} %param_3.933 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) %sub.73 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.933), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.64 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1340, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1577 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1577), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1570 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1570), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.60 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.1225 = f32[]{:T(128)} constant(0) - %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1225), dimensions={}, metadata={op_name="broadcast.109"} - %mul.1765 = f32[4,128,151936]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.769), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1765, %constant.1225), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %constant.1223 = f32[]{:T(128)} constant(0) + %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.109"} + %mul.1805 = f32[4,128,151936]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.769), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1805, %constant.1223), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_9.12 (reduce_sum.186: f32[], reduce_sum.190: f32[]) -> f32[] { @@ -203,15 +203,15 @@ StackFrames ROOT %reduce_sum.191 = f32[]{:T(128)} add(%reduce_sum.186, %reduce_sum.190), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.293 (param_0.1393: bf16[4,128,151936], param_1.1578: bf16[4,128]) -> f32[4,128] { - %param_0.1393 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1578 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1578), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.293 (param_0.1392: bf16[4,128,151936], param_1.1571: bf16[4,128]) -> f32[4,128] { + %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1571 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1571), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.70 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1346, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.54 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.1226 = f32[]{:T(128)} constant(0) - ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1226), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %constant.1224 = f32[]{:T(128)} constant(0) + ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1224), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_33.38 (reduce_sum.269: f32[], reduce_sum.270: f32[]) -> f32[] { @@ -220,12 +220,12 @@ StackFrames ROOT %reduce_sum.274 = f32[]{:T(128)} add(%reduce_sum.269, %reduce_sum.270), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.298 (param_0.1387: f32[4,6144,2048]) -> f32[] { - %param_0.1387 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) - %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.218 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%bitcast.347, %bitcast.347), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1219 = f32[]{:T(128)} constant(0) - ROOT %reduce.181 = f32[]{:T(128)} reduce(%square.218, %constant.1219), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.298 (param_0.1386: f32[4,6144,2048]) -> f32[] { + %param_0.1386 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) + %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1810 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%bitcast.347, %bitcast.347), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1217 = f32[]{:T(128)} constant(0) + ROOT %reduce.181 = f32[]{:T(128)} reduce(%mul.1810, %constant.1217), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_32.37 (reduce_sum.263: f32[], reduce_sum.267: f32[]) -> f32[] { @@ -240,35 +240,35 @@ StackFrames ROOT %reduce_sum.262 = f32[]{:T(128)} add(%reduce_sum.260, %reduce_sum.261), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.300 (param_0.1388: f32[4,2048,6144], param_1.1573: f32[4,2048,6144]) -> (f32[], f32[]) { - %param_0.1388 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) - %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.221 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.351, %bitcast.351), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1220 = f32[]{:T(128)} constant(0) - %reduce.182 = f32[]{:T(128)} reduce(%square.221, %constant.1220), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1573 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) - %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.224.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.355.clone.1, %bitcast.355.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.183.clone.1 = f32[]{:T(128)} reduce(%square.224.clone.1, %constant.1220), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.300 (param_0.1387: f32[4,2048,6144], param_1.1566: f32[4,2048,6144]) -> (f32[], f32[]) { + %param_0.1387 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) + %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1813 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.351, %bitcast.351), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1218 = f32[]{:T(128)} constant(0) + %reduce.182 = f32[]{:T(128)} reduce(%mul.1813, %constant.1218), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1566 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) + %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1566), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1816.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.355.clone.1, %bitcast.355.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.183.clone.1 = f32[]{:T(128)} reduce(%mul.1816.clone.1, %constant.1218), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.167 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.182, %reduce.183.clone.1) } -%fused_computation.303 (param_0.901: f32[6144,4,2048]) -> bf16[4,6144,2048] { - %param_0.901 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) - %copy.190 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.901), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.190), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.303 (param_0.900: f32[6144,4,2048]) -> bf16[4,6144,2048] { + %param_0.900 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) + %copy.192 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.900), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} + ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.304 (param_0.903: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.903 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.191 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.903), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.304 (param_0.902: f32[2048,4,6144]) -> bf16[4,2048,6144] { + %param_0.902 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.193 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.902), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} + ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.305 (param_0.905: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.905 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.192 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.905), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.305 (param_0.904: f32[2048,4,6144]) -> bf16[4,2048,6144] { + %param_0.904 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.194 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.904), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} + ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_62.67 (reduce_sum.416: f32[], reduce_sum.417: f32[]) -> f32[] { @@ -283,39 +283,39 @@ StackFrames ROOT %reduce_sum.340 = f32[]{:T(128)} add(%reduce_sum.338, %reduce_sum.339), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.306 (param_0.1377: f32[6144,4,2048], param_1.1565: f32[], param_2.1323: f32[], param_3.927: f32[], param_4.565: f32[6144,4,2048], param_5.477: f32[], param_6.367: f32[4,6144,2048], param_7.210: pred[], param_8.127: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { - %param_0.1377 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) +%fused_computation.306 (param_0.1376: f32[6144,4,2048], param_1.1558: f32[], param_2.1320: f32[], param_3.927: f32[], param_4.563: f32[6144,4,2048], param_5.476: f32[], param_6.365: f32[4,6144,2048], param_7.205: pred[], param_8.122: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { + %param_0.1376 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) %param_3.927 = f32[]{:T(128)S(6)} parameter(3) - %mul.1998.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_3.927), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.210 = pred[]{:T(512)S(6)} parameter(7) - %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.210), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.367 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) - %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.477 = f32[]{:T(128)} parameter(5) - %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2074.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_3.927), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.205 = pred[]{:T(512)S(6)} parameter(7) + %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.365 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) + %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.476 = f32[]{:T(128)} parameter(5) + %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.931.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%bitcast.482.clone.1, %div.932.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.303.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} select(%select_n.304.clone.1, %bitcast.482.clone.1, %div.931.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1146.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.886.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1146.clone.1), dimensions={}, metadata={op_name="broadcast.83"} - %mul.2004.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.127 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) + %mul.2080.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.122 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) %constant.1150.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2005.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1150.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2003.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.127, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.989.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2004.clone.1, %mul.2003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) - %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2081.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1150.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2079.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.122, %mul.2081.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.989.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2080.clone.1, %mul.2079.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) + %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.74.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %select_n.303.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1149.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2002.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1149.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2000.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%integer_pow.74.clone.1, %mul.2002.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.565 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) + %mul.2078.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1149.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2076.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%integer_pow.74.clone.1, %mul.2078.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.563 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) %constant.1148.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2001.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1148.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1999.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.565, %mul.2001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.988.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2000.clone.1, %mul.1999.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1565 = f32[]{:T(128)S(6)} parameter(1) - %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1565), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2077.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1148.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2075.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.563, %mul.2077.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.988.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2076.clone.1, %mul.2075.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) + %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.926.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.988.clone.1, %div.927.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.71.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} sqrt(%div.926.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1147.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -323,14 +323,14 @@ StackFrames %add.986.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%sqrt.71.clone.1, %add.987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.435.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%div.928.clone.1, %add.986.clone.1), metadata={op_name="multiply.52"} %div.925.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.989.clone.1, %multiply.435.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1997.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.985.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%div.925.clone.1, %mul.1997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1996.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%mul.1998.clone.1, %add.985.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1377, %mul.1996.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.225 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%add.984.clone.1, %add.984.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1209 = f32[]{:T(128)} constant(0) - %reduce.184 = f32[]{:T(128)} reduce(%square.225, %constant.1209), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2073.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1376, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.985.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%div.925.clone.1, %mul.2073.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2072.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%mul.2074.clone.1, %add.985.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1376, %mul.2072.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.176 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%add.984.clone.1, %add.984.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1207 = f32[]{:T(128)} constant(0) + %reduce.184 = f32[]{:T(128)} reduce(%square.176, %constant.1207), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1207), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.145 = (f32[]{:T(128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.184, %add.984.clone.1, %add.988.clone.1, %add.989.clone.1, %reduce.187.clone.1) } @@ -346,39 +346,39 @@ StackFrames ROOT %reduce_sum.337 = f32[]{:T(128)} add(%reduce_sum.332, %reduce_sum.333), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.307 (param_0.1378: f32[2048,4,6144], param_1.1566: f32[], param_2.1324: f32[], param_3.928: f32[], param_4.566: f32[2048,4,6144], param_5.478: f32[], param_6.368: f32[4,2048,6144], param_7.211: pred[], param_8.128: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) +%fused_computation.307 (param_0.1377: f32[2048,4,6144], param_1.1559: f32[], param_2.1321: f32[], param_3.928: f32[], param_4.564: f32[2048,4,6144], param_5.477: f32[], param_6.366: f32[4,2048,6144], param_7.206: pred[], param_8.123: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { + %param_0.1377 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) %param_3.928 = f32[]{:T(128)S(6)} parameter(3) - %mul.2008.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.928), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.211 = pred[]{:T(512)S(6)} parameter(7) - %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.211), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.368 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.478 = f32[]{:T(128)} parameter(5) - %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2084.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.928), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.206 = pred[]{:T(512)S(6)} parameter(7) + %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.366 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.477 = f32[]{:T(128)} parameter(5) + %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.939.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.484.clone.1, %div.940.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.307.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.308.clone.1, %bitcast.484.clone.1, %div.939.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1152.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.892.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1152.clone.1), dimensions={}, metadata={op_name="broadcast.85"} - %mul.2012.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.128 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) + %mul.2088.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.123 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) %constant.1156.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.891.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1156.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2011.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.128, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.994.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2012.clone.1, %mul.2011.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1324 = f32[]{:T(128)S(6)} parameter(2) - %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1324), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2087.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.123, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.994.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2088.clone.1, %mul.2087.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) + %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.75.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %select_n.307.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1155.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.890.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1155.clone.1), dimensions={}, metadata={op_name="broadcast.73"} - %mul.2010.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.75.clone.1, %broadcast.890.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.566 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) + %mul.2086.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.75.clone.1, %broadcast.890.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.564 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) %constant.1154.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.889.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1154.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2009.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.566, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.993.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2010.clone.1, %mul.2009.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1566 = f32[]{:T(128)S(6)} parameter(1) - %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1566), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2085.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.564, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.993.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2086.clone.1, %mul.2085.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) + %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.934.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.993.clone.1, %div.935.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.72.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.934.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1153.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -386,14 +386,14 @@ StackFrames %add.992.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.72.clone.1, %broadcast.887.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.436.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.936.clone.1, %add.992.clone.1), metadata={op_name="multiply.51"} %div.933.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.994.clone.1, %multiply.436.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2007.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.991.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.933.clone.1, %mul.2007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2006.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2008.clone.1, %add.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2006.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.226 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.990.clone.1, %add.990.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1210 = f32[]{:T(128)} constant(0) - %reduce.185 = f32[]{:T(128)} reduce(%square.226, %constant.1210), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1210), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2083.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.991.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.933.clone.1, %mul.2083.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2082.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2084.clone.1, %add.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1377, %mul.2082.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.177 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.990.clone.1, %add.990.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1208 = f32[]{:T(128)} constant(0) + %reduce.185 = f32[]{:T(128)} reduce(%square.177, %constant.1208), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1208), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.146 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.185, %add.990.clone.1, %add.993.clone.1, %add.994.clone.1, %reduce.188.clone.1) } @@ -409,39 +409,39 @@ StackFrames ROOT %reduce_sum.331 = f32[]{:T(128)} add(%reduce_sum.326, %reduce_sum.330), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.308 (param_0.1379: f32[2048,4,6144], param_1.1567: f32[], param_2.1325: f32[], param_3.929: f32[], param_4.567: f32[2048,4,6144], param_5.479: f32[], param_6.369: f32[4,2048,6144], param_7.212: pred[], param_8.129: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1379 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) +%fused_computation.308 (param_0.1378: f32[2048,4,6144], param_1.1560: f32[], param_2.1322: f32[], param_3.929: f32[], param_4.565: f32[2048,4,6144], param_5.478: f32[], param_6.367: f32[4,2048,6144], param_7.207: pred[], param_8.124: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { + %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) %param_3.929 = f32[]{:T(128)S(6)} parameter(3) - %mul.2015.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.929), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.212 = pred[]{:T(512)S(6)} parameter(7) - %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.212), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.369 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.369), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.479 = f32[]{:T(128)} parameter(5) - %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2091.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.929), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.207 = pred[]{:T(512)S(6)} parameter(7) + %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.367 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.478 = f32[]{:T(128)} parameter(5) + %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.947.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.486.clone.1, %div.948.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.311.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.312.clone.1, %bitcast.486.clone.1, %div.947.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1158.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.898.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1158.clone.1), dimensions={}, metadata={op_name="broadcast.85"} - %mul.2019.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.129 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) + %mul.2095.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.124 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) %constant.1162.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.897.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1162.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2018.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.129, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.999.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2019.clone.1, %mul.2018.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1325 = f32[]{:T(128)S(6)} parameter(2) - %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1325), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2094.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.124, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.999.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2095.clone.1, %mul.2094.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) + %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.76.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %select_n.311.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1161.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.896.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1161.clone.1), dimensions={}, metadata={op_name="broadcast.73"} - %mul.2017.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.76.clone.1, %broadcast.896.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.567 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) + %mul.2093.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.76.clone.1, %broadcast.896.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.565 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) %constant.1160.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.895.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1160.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2016.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.567, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.998.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2017.clone.1, %mul.2016.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1567 = f32[]{:T(128)S(6)} parameter(1) - %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1567), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2092.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.565, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.998.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2093.clone.1, %mul.2092.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) + %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.942.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.998.clone.1, %div.943.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.73.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.942.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1159.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -449,14 +449,14 @@ StackFrames %add.997.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.73.clone.1, %broadcast.893.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.437.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.944.clone.1, %add.997.clone.1), metadata={op_name="multiply.50"} %div.941.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.999.clone.1, %multiply.437.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2014.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1379, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.996.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.941.clone.1, %mul.2014.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2013.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2015.clone.1, %add.996.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1379, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.227 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.995.clone.1, %add.995.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1211 = f32[]{:T(128)} constant(0) - %reduce.186 = f32[]{:T(128)} reduce(%square.227, %constant.1211), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1211), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2090.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.996.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.941.clone.1, %mul.2090.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2089.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2091.clone.1, %add.996.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2089.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.178 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.995.clone.1, %add.995.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1209 = f32[]{:T(128)} constant(0) + %reduce.186 = f32[]{:T(128)} reduce(%square.178, %constant.1209), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.147 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.186, %add.995.clone.1, %add.998.clone.1, %add.999.clone.1, %reduce.189.clone.1) } @@ -466,12 +466,12 @@ StackFrames ROOT %reduce_sum.304 = f32[]{:T(128)} add(%reduce_sum.302, %reduce_sum.303), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.324 (param_0.1382: f32[4,2048,16,128]) -> f32[] { - %param_0.1382 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.230 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%bitcast.362, %bitcast.362), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1214 = f32[]{:T(128)} constant(0) - ROOT %reduce.190 = f32[]{:T(128)} reduce(%square.230, %constant.1214), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.324 (param_0.1381: f32[4,2048,16,128]) -> f32[] { + %param_0.1381 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1845 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%bitcast.362, %bitcast.362), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1212 = f32[]{:T(128)} constant(0) + ROOT %reduce.190 = f32[]{:T(128)} reduce(%mul.1845, %constant.1212), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_38.43 (reduce_sum.296: f32[], reduce_sum.297: f32[]) -> f32[] { @@ -480,18 +480,18 @@ StackFrames ROOT %reduce_sum.298 = f32[]{:T(128)} add(%reduce_sum.296, %reduce_sum.297), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.326 (param_0.1383: f32[4,16,128,2048]) -> f32[] { - %param_0.1383 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.233 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%bitcast.366, %bitcast.366), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1215 = f32[]{:T(128)} constant(0) - ROOT %reduce.191 = f32[]{:T(128)} reduce(%square.233, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.326 (param_0.1382: f32[4,16,128,2048]) -> f32[] { + %param_0.1382 = f32[4,16,128,2048]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1848 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%bitcast.366, %bitcast.366), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1213 = f32[]{:T(128)} constant(0) + ROOT %reduce.191 = f32[]{:T(128)} reduce(%mul.1848, %constant.1213), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.327 (param_0.950: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { - %param_0.950 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) - %copy.193 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.950), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.327 (param_0.949: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { + %param_0.949 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) + %copy.195 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.949), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} + ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.195), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_68.73 (reduce_sum.449: f32[], reduce_sum.450: f32[]) -> f32[] { @@ -506,39 +506,39 @@ StackFrames ROOT %reduce_sum.373 = f32[]{:T(128)} add(%reduce_sum.368, %reduce_sum.372), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.328 (param_0.1371: f32[2048,4,16,128], param_1.1559: f32[], param_2.1317: f32[], param_3.921: f32[], param_4.559: f32[2048,4,16,128], param_5.471: f32[], param_6.361: f32[4,2048,16,128], param_7.204: pred[], param_8.121: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { - %param_0.1371 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.328 (param_0.1370: f32[2048,4,16,128], param_1.1552: f32[], param_2.1314: f32[], param_3.921: f32[], param_4.557: f32[2048,4,16,128], param_5.470: f32[], param_6.359: f32[4,2048,16,128], param_7.199: pred[], param_8.116: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { + %param_0.1370 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.921 = f32[]{:T(128)S(6)} parameter(3) - %mul.1950.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_3.921), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.204 = pred[]{:T(512)S(6)} parameter(7) - %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.361 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.471 = f32[]{:T(128)} parameter(5) - %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2026.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_3.921), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.199 = pred[]{:T(512)S(6)} parameter(7) + %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.199), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.359 = f32[4,2048,16,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.470 = f32[]{:T(128)} parameter(5) + %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.883.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%bitcast.470.clone.1, %div.884.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.279.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} select(%select_n.280.clone.1, %bitcast.470.clone.1, %div.883.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1110.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.858.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1110.clone.1), dimensions={}, metadata={op_name="broadcast.75"} - %mul.1956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.121 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.2032.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.116 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1114.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1114.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.121, %mul.1957.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1956.clone.1, %mul.1955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) - %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2033.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1114.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2031.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.116, %mul.2033.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.2032.clone.1, %mul.2031.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) + %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.68.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %select_n.279.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1113.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1113.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.68.clone.1, %mul.1954.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.559 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.2030.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1113.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2028.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.68.clone.1, %mul.2030.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.557 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1112.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1112.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1951.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.559, %mul.1953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1952.clone.1, %mul.1951.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) - %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2029.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1112.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2027.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.2028.clone.1, %mul.2027.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1552 = f32[]{:T(128)S(6)} parameter(1) + %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1552), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.878.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.956.clone.1, %div.879.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.65.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} sqrt(%div.878.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1111.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -546,14 +546,14 @@ StackFrames %add.954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%sqrt.65.clone.1, %add.955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.429.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%div.880.clone.1, %add.954.clone.1), metadata={op_name="multiply.58"} %div.877.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.957.clone.1, %multiply.429.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1949.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%div.877.clone.1, %mul.1949.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1948.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%mul.1950.clone.1, %add.953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.1948.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.234 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%add.952.clone.1, %add.952.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1203 = f32[]{:T(128)} constant(0) - %reduce.192 = f32[]{:T(128)} reduce(%square.234, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2025.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1370, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%div.877.clone.1, %mul.2025.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2024.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%mul.2026.clone.1, %add.953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1370, %mul.2024.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.179 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%add.952.clone.1, %add.952.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1201 = f32[]{:T(128)} constant(0) + %reduce.192 = f32[]{:T(128)} reduce(%square.179, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.148 = (f32[]{:T(128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.192, %add.952.clone.1, %add.956.clone.1, %add.957.clone.1, %reduce.194.clone.1) } @@ -569,39 +569,39 @@ StackFrames ROOT %reduce_sum.367 = f32[]{:T(128)} add(%reduce_sum.365, %reduce_sum.366), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.329 (param_0.1372: f32[16,4,128,2048], param_1.1560: f32[], param_2.1318: f32[], param_3.922: f32[], param_4.560: f32[16,4,128,2048], param_5.472: f32[], param_6.362: f32[4,16,128,2048], param_7.205: pred[], param_8.122: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { - %param_0.1372 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.329 (param_0.1371: f32[16,4,128,2048], param_1.1553: f32[], param_2.1315: f32[], param_3.922: f32[], param_4.558: f32[16,4,128,2048], param_5.471: f32[], param_6.360: f32[4,16,128,2048], param_7.200: pred[], param_8.117: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { + %param_0.1371 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) %param_3.922 = f32[]{:T(128)S(6)} parameter(3) - %mul.1960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_3.922), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.205 = pred[]{:T(512)S(6)} parameter(7) - %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.362 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.472 = f32[]{:T(128)} parameter(5) - %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2036.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_3.922), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.200 = pred[]{:T(512)S(6)} parameter(7) + %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.200), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.360 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.471 = f32[]{:T(128)} parameter(5) + %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.891.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%bitcast.472.clone.1, %div.892.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.283.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} select(%select_n.284.clone.1, %bitcast.472.clone.1, %div.891.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1116.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.860.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1116.clone.1), dimensions={}, metadata={op_name="broadcast.76"} - %mul.1966.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.122 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) + %mul.2042.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.117 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) %constant.1120.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1967.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1120.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1965.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.122, %mul.1967.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1966.clone.1, %mul.1965.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) - %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2043.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1120.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2041.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.117, %mul.2043.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.2042.clone.1, %mul.2041.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) + %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.69.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %select_n.283.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1119.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1964.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1119.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%integer_pow.69.clone.1, %mul.1964.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.560 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) + %mul.2040.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1119.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2038.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%integer_pow.69.clone.1, %mul.2040.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.558 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) %constant.1118.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1118.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.560, %mul.1963.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1962.clone.1, %mul.1961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) - %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2039.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1118.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2037.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.558, %mul.2039.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.2038.clone.1, %mul.2037.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1553 = f32[]{:T(128)S(6)} parameter(1) + %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1553), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.886.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.962.clone.1, %div.887.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.66.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} sqrt(%div.886.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1117.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -609,14 +609,14 @@ StackFrames %add.960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%sqrt.66.clone.1, %add.961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.430.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%div.888.clone.1, %add.960.clone.1), metadata={op_name="multiply.57"} %div.885.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.963.clone.1, %multiply.430.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1372, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%div.885.clone.1, %mul.1959.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%mul.1960.clone.1, %add.959.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1372, %mul.1958.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.235 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%add.958.clone.1, %add.958.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1204 = f32[]{:T(128)} constant(0) - %reduce.193 = f32[]{:T(128)} reduce(%square.235, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2035.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%div.885.clone.1, %mul.2035.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2034.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%mul.2036.clone.1, %add.959.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.2034.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.180 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%add.958.clone.1, %add.958.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1202 = f32[]{:T(128)} constant(0) + %reduce.193 = f32[]{:T(128)} reduce(%square.180, %constant.1202), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1202), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.149 = (f32[]{:T(128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.193, %add.958.clone.1, %add.962.clone.1, %add.963.clone.1, %reduce.195.clone.1) } @@ -632,23 +632,23 @@ StackFrames ROOT %reduce_sum.289 = f32[]{:T(128)} add(%reduce_sum.284, %reduce_sum.288), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.341 (param_0.1385: f32[4,2048,8,128], param_1.1571: f32[4,2048,8,128]) -> (f32[], f32[]) { - %param_0.1385 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) - %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.238 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1217 = f32[]{:T(128)} constant(0) - %reduce.196 = f32[]{:T(128)} reduce(%square.238, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1571 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(1) - %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.241.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.197.clone.1 = f32[]{:T(128)} reduce(%square.241.clone.1, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.341 (param_0.1384: f32[4,2048,8,128], param_1.1564: f32[4,2048,8,128]) -> (f32[], f32[]) { + %param_0.1384 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1863 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1215 = f32[]{:T(128)} constant(0) + %reduce.196 = f32[]{:T(128)} reduce(%mul.1863, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1564 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) + %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1866.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.197.clone.1 = f32[]{:T(128)} reduce(%mul.1866.clone.1, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.168 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.196, %reduce.197.clone.1) } -%fused_computation.344 (param_0.982: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.982 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %copy.194 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.982), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.344 (param_0.981: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { + %param_0.981 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.196 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.981), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} + ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.196), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_70.75 (reduce_sum.458: f32[], reduce_sum.459: f32[]) -> f32[] { @@ -663,39 +663,39 @@ StackFrames ROOT %reduce_sum.382 = f32[]{:T(128)} add(%reduce_sum.380, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.345 (param_0.1369: f32[2048,4,8,128], param_1.1557: f32[], param_2.1315: f32[], param_3.919: f32[], param_4.557: f32[2048,4,8,128], param_5.469: f32[], param_6.359: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1369 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.345 (param_0.1368: f32[2048,4,8,128], param_1.1550: f32[], param_2.1312: f32[], param_3.919: f32[], param_4.555: f32[2048,4,8,128], param_5.468: f32[], param_6.357: f32[4,2048,8,128], param_7.197: pred[], param_8.114: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { + %param_0.1368 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.919 = f32[]{:T(128)S(6)} parameter(3) - %mul.1936.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.919), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.202 = pred[]{:T(512)S(6)} parameter(7) - %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.359 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.469 = f32[]{:T(128)} parameter(5) - %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2012.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.919), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.197 = pred[]{:T(512)S(6)} parameter(7) + %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.357 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.357), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.468 = f32[]{:T(128)} parameter(5) + %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.867.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.466.clone.1, %div.868.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.271.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.272.clone.1, %bitcast.466.clone.1, %div.867.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1098.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.850.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1098.clone.1), dimensions={}, metadata={op_name="broadcast.80"} - %mul.1940.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.2016.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.114 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1102.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.849.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1102.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.1939.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.946.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1940.clone.1, %mul.1939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) - %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2015.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.114, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.946.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2016.clone.1, %mul.2015.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1312 = f32[]{:T(128)S(6)} parameter(2) + %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1312), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.66.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %select_n.271.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1101.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.848.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1101.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1938.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.66.clone.1, %broadcast.848.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.557 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.2014.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.66.clone.1, %broadcast.848.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.555 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1100.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.847.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1100.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1937.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.945.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1938.clone.1, %mul.1937.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) - %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2013.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.555, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.945.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2014.clone.1, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1550 = f32[]{:T(128)S(6)} parameter(1) + %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.862.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.945.clone.1, %div.863.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.63.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.862.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1099.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -703,14 +703,14 @@ StackFrames %add.944.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.845.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.427.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.864.clone.1, %add.944.clone.1), metadata={op_name="multiply.60"} %div.861.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.946.clone.1, %multiply.427.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1935.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1369, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.943.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.861.clone.1, %mul.1935.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1934.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1936.clone.1, %add.943.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1369, %mul.1934.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.242 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.942.clone.1, %add.942.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1201 = f32[]{:T(128)} constant(0) - %reduce.198 = f32[]{:T(128)} reduce(%square.242, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2011.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1368, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.943.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.861.clone.1, %mul.2011.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2010.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.2012.clone.1, %add.943.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1368, %mul.2010.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.181 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.942.clone.1, %add.942.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1199 = f32[]{:T(128)} constant(0) + %reduce.198 = f32[]{:T(128)} reduce(%square.181, %constant.1199), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1199), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.150 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.198, %add.942.clone.1, %add.945.clone.1, %add.946.clone.1, %reduce.200.clone.1) } @@ -726,39 +726,39 @@ StackFrames ROOT %reduce_sum.358 = f32[]{:T(128)} add(%reduce_sum.353, %reduce_sum.354), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.346 (param_0.1374: f32[2048,4,8,128], param_1.1562: f32[], param_2.1320: f32[], param_3.924: f32[], param_4.562: f32[2048,4,8,128], param_5.474: f32[], param_6.364: f32[4,2048,8,128], param_7.207: pred[], param_8.124: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1374 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.346 (param_0.1373: f32[2048,4,8,128], param_1.1555: f32[], param_2.1317: f32[], param_3.924: f32[], param_4.560: f32[2048,4,8,128], param_5.473: f32[], param_6.362: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { + %param_0.1373 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) %param_3.924 = f32[]{:T(128)S(6)} parameter(3) - %mul.1977.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.924), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.207 = pred[]{:T(512)S(6)} parameter(7) - %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.364 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.474 = f32[]{:T(128)} parameter(5) - %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2053.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.924), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.202 = pred[]{:T(512)S(6)} parameter(7) + %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.362 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.473 = f32[]{:T(128)} parameter(5) + %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.907.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.476.clone.1, %div.908.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.291.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.292.clone.1, %bitcast.476.clone.1, %div.907.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1128.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.872.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1128.clone.1), dimensions={}, metadata={op_name="broadcast.80"} - %mul.1981.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.124 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.2057.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1132.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.871.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1132.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.1980.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.124, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.973.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1981.clone.1, %mul.1980.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) - %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2056.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.973.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2057.clone.1, %mul.2056.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) + %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.71.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %select_n.291.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1131.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.870.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1131.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1979.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.71.clone.1, %broadcast.870.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.562 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.2055.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.71.clone.1, %broadcast.870.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.560 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1130.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.869.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1130.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1978.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.562, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.972.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1979.clone.1, %mul.1978.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1562 = f32[]{:T(128)S(6)} parameter(1) - %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1562), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2054.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.560, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.972.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2055.clone.1, %mul.2054.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1555 = f32[]{:T(128)S(6)} parameter(1) + %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1555), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.902.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.972.clone.1, %div.903.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.68.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.902.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1129.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -766,32 +766,32 @@ StackFrames %add.971.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.68.clone.1, %broadcast.867.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.432.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.904.clone.1, %add.971.clone.1), metadata={op_name="multiply.55"} %div.901.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.973.clone.1, %multiply.432.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1976.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1374, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.970.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.901.clone.1, %mul.1976.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1975.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1977.clone.1, %add.970.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1374, %mul.1975.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.243 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.969.clone.1, %add.969.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1206 = f32[]{:T(128)} constant(0) - %reduce.199 = f32[]{:T(128)} reduce(%square.243, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) + %mul.2052.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1373, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.970.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.901.clone.1, %mul.2052.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2051.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.2053.clone.1, %add.970.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1373, %mul.2051.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.182 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.969.clone.1, %add.969.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1204 = f32[]{:T(128)} constant(0) + %reduce.199 = f32[]{:T(128)} reduce(%square.182, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) } -%fused_computation.362 (param_0.1056: bf16[4,128,2048], param_1.1117: f32[4,128], param_2.830: f32[4,128], param_3.495: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { - %param_3.495 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) +%fused_computation.362 (param_0.1055: bf16[4,128,2048], param_1.1114: f32[4,128], param_2.829: f32[4,128], param_3.497: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { + %param_3.497 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) %param_4.296 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.448 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.438 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.495, %dot_general.448), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.438), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.830 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1851 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.830), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1843 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1363, %mul.1851), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.1056 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1056), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1117 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1850 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1117), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1849 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1374, %mul.1850), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %add_any.193 = f32[4,128,2048]{2,1,0:T(8,128)} add(%mul.1843, %mul.1849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} + %dot_general.451 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.441 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.497, %dot_general.451), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.441), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_2.829 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %mul.1912 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.829), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1904 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1363, %mul.1912), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_0.1055 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1055), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1911 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1114), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1910 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1374, %mul.1911), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %add_any.193 = f32[4,128,2048]{2,1,0:T(8,128)} add(%mul.1904, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} ROOT %convert_element_type.1361 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.193), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} } @@ -801,12 +801,12 @@ StackFrames ROOT %reduce_sum.185 = f32[]{:T(128)} add(%reduce_sum.171, %reduce_sum.184), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.363 (param_0.1394: bf16[4,128,2048]) -> f32[4,128] { - %param_0.1394 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1394), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.246 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1365, %convert_element_type.1365), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} - %constant.1227 = f32[]{:T(128)} constant(0) - ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.246, %constant.1227), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} +%fused_computation.363 (param_0.1393: bf16[4,128,2048]) -> f32[4,128] { + %param_0.1393 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.185 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1365, %convert_element_type.1365), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} + %constant.1225 = f32[]{:T(128)} constant(0) + ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.185, %constant.1225), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_12.15 (reduce_sum.198: f32[], reduce_sum.199: f32[]) -> f32[] { @@ -815,17 +815,17 @@ StackFrames ROOT %reduce_sum.200 = f32[]{:T(128)} add(%reduce_sum.198, %reduce_sum.199), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.365 (param_0.1389: bf16[4,128,2048], param_1.1574: bf16[4,128,2048], param_2.1328: bf16[2048]) -> f32[4,128] { - %param_0.1389 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1389), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1328 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.447 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1328), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1574, %dot_general.447), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.437), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1847 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1372, %convert_element_type.1371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %constant.1221 = f32[]{:T(128)} constant(0) - ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1847, %constant.1221), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} +%fused_computation.365 (param_0.1388: bf16[4,128,2048], param_1.1567: bf16[4,128,2048], param_2.1325: bf16[2048]) -> f32[4,128] { + %param_0.1388 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1388), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1325 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1325), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.440 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1567, %dot_general.450), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.440), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.1908 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1372, %convert_element_type.1371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %constant.1219 = f32[]{:T(128)} constant(0) + ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1908, %constant.1219), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_10.13 (dot_general.190: bf16[], dot_general.191: bf16[]) -> bf16[] { @@ -834,64 +834,64 @@ StackFrames ROOT %add.419 = bf16[]{:T(256)} add(%dot_general.190, %dot_general.191), metadata={op_name="add.82"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.285.clone.clone (param_0.1351: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1351 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%fused_computation.285.clone.clone (param_0.1350: bf16[151936,2048]) -> bf16[151936,2048,1] { + %param_0.1350 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.530 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1350), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} } -%fused_computation.289.clone.1.clone.clone (param_0.1352: bf16[4,128,151936], param_1.1546: s32[4,128], param_2.1285: f32[4,128], param_3.906: f32[4,128], param_4.542: bf16[4,128], param_5.442: f32[4,128]) -> bf16[4,128,151936] { - %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2075 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.442), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.289.clone.1.clone.clone (param_0.1351: bf16[4,128,151936], param_1.1539: s32[4,128], param_2.1282: f32[4,128], param_3.906: f32[4,128], param_4.540: bf16[4,128], param_5.441: f32[4,128]) -> bf16[4,128,151936] { + %param_5.441 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.2143 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.441), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.906 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2074 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1352 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1444 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1352), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.542 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.542), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1444, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %mul.2142 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1351 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1438 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.540 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.540), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1438, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.60 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2073 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2074, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1285 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1285), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2073, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1546 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.2141 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2142, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1282 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1282), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2141, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1539 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.42 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.41 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1443 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1443), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2072 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2075, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1442 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2072), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convert_element_type.1437 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1437), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2140 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2143, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1436 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2140), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} } -%fused_computation.366 (param_0.1350: f32[4,128], param_1.1545: bf16[4,128,2048], param_2.1286: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.543: s32[4,128], param_5.443: f32[4,128], param_6.340: f32[4,128], param_7.199: bf16[4,128], param_8.116: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { +%fused_computation.366 (param_0.1349: f32[4,128], param_1.1538: bf16[4,128,2048], param_2.1283: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.541: s32[4,128], param_5.442: f32[4,128], param_6.338: f32[4,128], param_7.194: bf16[4,128], param_8.111: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { %param_3.907 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.443 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.340 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.199 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.116 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.543, %param_5.443, %param_6.340, %param_7.199, /*index=5*/%param_8.116), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1286 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) - %fusion.251.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1286), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.251.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %param_1.1545 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1545), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1350 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1862 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1350), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1861 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1384, %mul.1862), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1383 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1861), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_4.541 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.338 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.194 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) + %param_8.111 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) + %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.541, %param_5.442, %param_6.338, %param_7.194, /*index=5*/%param_8.111), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_2.1283 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) + %fusion.250.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1283), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.250.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %param_1.1538 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1349 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1923 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1349), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1922 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1384, %mul.1923), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1383 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1922), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %multiply.420 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%convolution.84.clone.1, %convert_element_type.1383), metadata={op_name="multiply.362"} %constant.1050 = bf16[]{:T(256)} constant(0) %reduce.204 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%multiply.420, %constant.1050), dimensions={0,1}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} ROOT %tuple.165 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.204, %convolution.84.clone.1) } -%fused_computation.374 (param_0.1088: f32[64], param_1.1150: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1150 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1150), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.1088 = f32[64]{0:T(128)S(1)} parameter(0) - %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1088), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} +%fused_computation.374 (param_0.1087: f32[64], param_1.1147: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1147 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1147), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %param_0.1087 = f32[64]{0:T(128)S(1)} parameter(0) + %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1087), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %div.717 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.720, %div.718), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.717), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} %convert_element_type.1392 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} @@ -900,19 +900,19 @@ StackFrames ROOT %tuple.158 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1392, %convert_element_type.1391.clone.1) } -%fused_computation.375 (param_0.1085: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1085 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.375 (param_0.1084: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.1084 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1042 = bf16[]{:T(256)} constant(-inf) - %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1084, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1084, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.42 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.46, %pad.45), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } -%fused_computation.376 (param_0.1087: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1087 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.376 (param_0.1086: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.1086 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1041 = bf16[]{:T(256)} constant(-inf) - %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1086, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1086, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.43 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.48, %pad.47), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } @@ -928,16 +928,16 @@ StackFrames ROOT %reduce_sum.277 = f32[]{:T(128)} add(%reduce_sum.275, %reduce_sum.276), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.380 (param_0.1386: f32[4,2048], param_1.1572: f32[4,2048]) -> (f32[], f32[]) { - %param_0.1386 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.249 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.404, %bitcast.404), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1218 = f32[]{:T(128)} constant(0) - %reduce.205 = f32[]{:T(128)} reduce(%square.249, %constant.1218), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1572 = f32[4,2048]{1,0:T(4,128)} parameter(1) - %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.252.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.408.clone.1, %bitcast.408.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.206.clone.1 = f32[]{:T(128)} reduce(%square.252.clone.1, %constant.1218), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.380 (param_0.1385: f32[4,2048], param_1.1565: f32[4,2048]) -> (f32[], f32[]) { + %param_0.1385 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1932 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.404, %bitcast.404), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1216 = f32[]{:T(128)} constant(0) + %reduce.205 = f32[]{:T(128)} reduce(%mul.1932, %constant.1216), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1565 = f32[4,2048]{1,0:T(4,128)} parameter(1) + %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1935.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.408.clone.1, %bitcast.408.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.206.clone.1 = f32[]{:T(128)} reduce(%mul.1935.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.169 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.205, %reduce.206.clone.1) } @@ -953,39 +953,39 @@ StackFrames ROOT %reduce_sum.352 = f32[]{:T(128)} add(%reduce_sum.347, %reduce_sum.351), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.383 (param_0.1375: f32[2048,4], param_1.1563: f32[], param_2.1321: f32[], param_3.925: f32[], param_4.563: f32[2048,4], param_5.475: f32[], param_6.365: f32[4,2048], param_7.208: pred[], param_8.125: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.383 (param_0.1374: f32[2048,4], param_1.1556: f32[], param_2.1318: f32[], param_3.925: f32[], param_4.561: f32[2048,4], param_5.474: f32[], param_6.363: f32[4,2048], param_7.203: pred[], param_8.120: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1374 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.925 = f32[]{:T(128)S(6)} parameter(3) - %mul.1984.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.925), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.208 = pred[]{:T(512)S(6)} parameter(7) - %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.365 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.475 = f32[]{:T(128)} parameter(5) - %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2060.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.925), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.203 = pred[]{:T(512)S(6)} parameter(7) + %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.363 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.474 = f32[]{:T(128)} parameter(5) + %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.915.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.478.clone.1, %div.916.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.295.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.296.clone.1, %bitcast.478.clone.1, %div.915.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1134.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.878.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1134.clone.1), dimensions={}, metadata={op_name="broadcast.82"} - %mul.1988.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.125 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.2064.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.120 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1138.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.877.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1138.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.1987.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.125, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.978.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1988.clone.1, %mul.1987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) - %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2063.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.978.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2064.clone.1, %mul.2063.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) + %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.72.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %select_n.295.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1137.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.876.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1137.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1986.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.72.clone.1, %broadcast.876.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.563 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.2062.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.72.clone.1, %broadcast.876.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.561 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1136.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.875.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1136.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1985.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.563, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.977.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1986.clone.1, %mul.1985.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1563 = f32[]{:T(128)S(6)} parameter(1) - %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1563), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2061.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.977.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2062.clone.1, %mul.2061.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) + %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.910.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.977.clone.1, %div.911.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.69.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.910.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1135.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -993,14 +993,14 @@ StackFrames %add.976.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.69.clone.1, %broadcast.873.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.433.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.912.clone.1, %add.976.clone.1), metadata={op_name="multiply.54"} %div.909.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.978.clone.1, %multiply.433.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1983.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.975.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.909.clone.1, %mul.1983.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1982.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1984.clone.1, %add.975.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.1982.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.253 = f32[2048,4]{0,1:T(4,128)} multiply(%add.974.clone.1, %add.974.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1207 = f32[]{:T(128)} constant(0) - %reduce.207 = f32[]{:T(128)} reduce(%square.253, %constant.1207), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1207), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2059.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1374, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.975.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.909.clone.1, %mul.2059.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2058.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.2060.clone.1, %add.975.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1374, %mul.2058.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.186 = f32[2048,4]{0,1:T(4,128)} multiply(%add.974.clone.1, %add.974.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1205 = f32[]{:T(128)} constant(0) + %reduce.207 = f32[]{:T(128)} reduce(%square.186, %constant.1205), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.152 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.207, %add.974.clone.1, %add.977.clone.1, %add.978.clone.1, %reduce.209.clone.1) } @@ -1016,39 +1016,39 @@ StackFrames ROOT %reduce_sum.346 = f32[]{:T(128)} add(%reduce_sum.344, %reduce_sum.345), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.384 (param_0.1376: f32[2048,4], param_1.1564: f32[], param_2.1322: f32[], param_3.926: f32[], param_4.564: f32[2048,4], param_5.476: f32[], param_6.366: f32[4,2048], param_7.209: pred[], param_8.126: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1376 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.384 (param_0.1375: f32[2048,4], param_1.1557: f32[], param_2.1319: f32[], param_3.926: f32[], param_4.562: f32[2048,4], param_5.475: f32[], param_6.364: f32[4,2048], param_7.204: pred[], param_8.121: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.926 = f32[]{:T(128)S(6)} parameter(3) - %mul.1991.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.926), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.209 = pred[]{:T(512)S(6)} parameter(7) - %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.209), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.366 = f32[4,2048]{1,0:T(4,128)} parameter(6) - %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.476 = f32[]{:T(128)} parameter(5) - %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2067.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.926), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.204 = pred[]{:T(512)S(6)} parameter(7) + %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.364 = f32[4,2048]{1,0:T(4,128)} parameter(6) + %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.475 = f32[]{:T(128)} parameter(5) + %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.923.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.480.clone.1, %div.924.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.299.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.300.clone.1, %bitcast.480.clone.1, %div.923.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1140.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.884.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1140.clone.1), dimensions={}, metadata={op_name="broadcast.82"} - %mul.1995.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.126 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.2071.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.121 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1144.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.883.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1144.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.1994.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.126, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.983.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1995.clone.1, %mul.1994.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) - %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2070.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.121, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.983.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2071.clone.1, %mul.2070.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) + %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.73.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %select_n.299.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1143.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.882.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1143.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1993.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.73.clone.1, %broadcast.882.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.564 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.2069.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.73.clone.1, %broadcast.882.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.562 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1142.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.881.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1142.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1992.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.564, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.982.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1993.clone.1, %mul.1992.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1564 = f32[]{:T(128)S(6)} parameter(1) - %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1564), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2068.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.562, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.982.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2069.clone.1, %mul.2068.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) + %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.918.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.982.clone.1, %div.919.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.70.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.918.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1141.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1056,14 +1056,14 @@ StackFrames %add.981.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.70.clone.1, %broadcast.879.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.434.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.920.clone.1, %add.981.clone.1), metadata={op_name="multiply.53"} %div.917.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.983.clone.1, %multiply.434.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1990.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1376, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.980.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.917.clone.1, %mul.1990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1989.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1991.clone.1, %add.980.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1376, %mul.1989.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.254 = f32[2048,4]{0,1:T(4,128)} multiply(%add.979.clone.1, %add.979.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1208 = f32[]{:T(128)} constant(0) - %reduce.208 = f32[]{:T(128)} reduce(%square.254, %constant.1208), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1208), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2066.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.980.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.917.clone.1, %mul.2066.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2065.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.2067.clone.1, %add.980.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.2065.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.187 = f32[2048,4]{0,1:T(4,128)} multiply(%add.979.clone.1, %add.979.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1206 = f32[]{:T(128)} constant(0) + %reduce.208 = f32[]{:T(128)} reduce(%square.187, %constant.1206), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1206), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.153 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.208, %add.979.clone.1, %add.982.clone.1, %add.983.clone.1, %reduce.210.clone.1) } @@ -1073,12 +1073,12 @@ StackFrames ROOT %reduce_sum.197 = f32[]{:T(128)} add(%reduce_sum.192, %reduce_sum.193), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.395 (param_0.1390: bf16[2048]) -> f32[] { - %param_0.1390 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.257 = f32[2048]{0:T(1024)} multiply(%convert_element_type.1396, %convert_element_type.1396), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1222 = f32[]{:T(128)} constant(0) - ROOT %reduce.211 = f32[]{:T(128)} reduce(%square.257, %constant.1222), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.395 (param_0.1389: bf16[2048]) -> f32[] { + %param_0.1389 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.1952 = f32[2048]{0:T(1024)} multiply(%convert_element_type.1396, %convert_element_type.1396), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1220 = f32[]{:T(128)} constant(0) + ROOT %reduce.211 = f32[]{:T(128)} reduce(%mul.1952, %constant.1220), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_59.64 (reduce_sum.401: f32[], reduce_sum.402: f32[]) -> f32[] { @@ -1093,39 +1093,39 @@ StackFrames ROOT %reduce_sum.325 = f32[]{:T(128)} add(%reduce_sum.323, %reduce_sum.324), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.396 (param_0.1380: f32[2048], param_1.1568: f32[], param_2.1326: f32[], param_3.930: f32[], param_4.568: f32[2048], param_5.480: f32[], param_6.370: bf16[2048], param_7.213: pred[], param_8.130: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { - %param_0.1380 = f32[2048]{0:T(1024)S(1)} parameter(0) +%fused_computation.396 (param_0.1379: f32[2048], param_1.1561: f32[], param_2.1323: f32[], param_3.930: f32[], param_4.566: f32[2048], param_5.479: f32[], param_6.368: bf16[2048], param_7.208: pred[], param_8.125: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { + %param_0.1379 = f32[2048]{0:T(1024)S(1)} parameter(0) %param_3.930 = f32[]{:T(128)S(6)} parameter(3) - %mul.2022.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_3.930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.213 = pred[]{:T(512)S(6)} parameter(7) - %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.213), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.370 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.480 = f32[]{:T(128)} parameter(5) - %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.480), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2098.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_3.930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.208 = pred[]{:T(512)S(6)} parameter(7) + %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.368 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_5.479 = f32[]{:T(128)} parameter(5) + %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.955.clone.1 = f32[2048]{0:T(1024)} divide(%convert_element_type.1411.clone.1, %div.956.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.315.clone.1 = f32[2048]{0:T(1024)} select(%select_n.316.clone.1, %convert_element_type.1411.clone.1, %div.955.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1164.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.900.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1164.clone.1), dimensions={}, metadata={op_name="broadcast.86"} - %mul.2028.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.130 = f32[2048]{0:T(1024)S(1)} parameter(8) + %mul.2104.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.125 = f32[2048]{0:T(1024)S(1)} parameter(8) %constant.1168.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2029.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1168.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2027.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.130, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1005.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2028.clone.1, %mul.2027.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1326 = f32[]{:T(128)S(6)} parameter(2) - %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1326), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2105.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1168.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2103.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.125, %mul.2105.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1005.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2104.clone.1, %mul.2103.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) + %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.77.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %select_n.315.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1167.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2026.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1167.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2024.clone.1 = f32[2048]{0:T(1024)} multiply(%integer_pow.77.clone.1, %mul.2026.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.568 = f32[2048]{0:T(1024)S(1)} parameter(4) + %mul.2102.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1167.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2100.clone.1 = f32[2048]{0:T(1024)} multiply(%integer_pow.77.clone.1, %mul.2102.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.566 = f32[2048]{0:T(1024)S(1)} parameter(4) %constant.1166.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2025.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1166.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2023.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.568, %mul.2025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1004.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2024.clone.1, %mul.2023.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) - %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2101.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1166.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2099.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.566, %mul.2101.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1004.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2100.clone.1, %mul.2099.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) + %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.950.clone.1 = f32[2048]{0:T(1024)} divide(%add.1004.clone.1, %div.951.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.74.clone.1 = f32[2048]{0:T(1024)} sqrt(%div.950.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1165.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1133,37 +1133,37 @@ StackFrames %add.1002.clone.1 = f32[2048]{0:T(1024)} add(%sqrt.74.clone.1, %add.1003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.438.clone.1 = f32[2048]{0:T(1024)} multiply(%div.952.clone.1, %add.1002.clone.1), metadata={op_name="multiply.49"} %div.949.clone.1 = f32[2048]{0:T(1024)} divide(%add.1005.clone.1, %multiply.438.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2021.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1380, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1001.clone.1 = f32[2048]{0:T(1024)} add(%div.949.clone.1, %mul.2021.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2020.clone.1 = f32[2048]{0:T(1024)} multiply(%mul.2022.clone.1, %add.1001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1380, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.258 = f32[2048]{0:T(1024)} multiply(%add.1000.clone.1, %add.1000.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1212 = f32[]{:T(128)} constant(0) - %reduce.212 = f32[]{:T(128)} reduce(%square.258, %constant.1212), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1212), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2097.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1379, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1001.clone.1 = f32[2048]{0:T(1024)} add(%div.949.clone.1, %mul.2097.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2096.clone.1 = f32[2048]{0:T(1024)} multiply(%mul.2098.clone.1, %add.1001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1379, %mul.2096.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.188 = f32[2048]{0:T(1024)} multiply(%add.1000.clone.1, %add.1000.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1210 = f32[]{:T(128)} constant(0) + %reduce.212 = f32[]{:T(128)} reduce(%square.188, %constant.1210), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1210), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.156 = (f32[]{:T(128)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.212, %add.1000.clone.1, %add.1004.clone.1, %add.1005.clone.1, %reduce.213.clone.1) } -%fused_computation.402 (param_0.1150: s32[512]) -> s32[1024] { +%fused_computation.402 (param_0.1149: s32[512]) -> s32[1024] { %constant.972 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.815 = s32[1024]{0:T(1024)} broadcast(%constant.972), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.1150 = s32[512]{0:T(512)S(1)} parameter(0) + %param_0.1149 = s32[512]{0:T(512)S(1)} parameter(0) %constant.973 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1150, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1149, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %constant.971 = s32[] constant(151935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.814 = s32[1024]{0:T(1024)} broadcast(%constant.971), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.815, %pad.49, %broadcast.814), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } -%fused_computation.405 (param_0.1149: s32[4,128]) -> s32[512] { - %param_0.1149 = s32[4,128]{1,0:T(4,128)} parameter(0) +%fused_computation.405 (param_0.1148: s32[4,128]) -> s32[512] { + %param_0.1148 = s32[4,128]{1,0:T(4,128)} parameter(0) %constant.1065 = s32[]{:T(128)} constant(0) %broadcast.834 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1065), dimensions={}, metadata={op_name="broadcast.95"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1149, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} + %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1148, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} %constant.1051 = s32[]{:T(128)} constant(151936) %add.925 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1051), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1149, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1149), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} + %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1148, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} + %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1148), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} ROOT %bitcast.409 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} } @@ -1179,16 +1179,16 @@ StackFrames ROOT %reduce_sum.295 = f32[]{:T(128)} add(%reduce_sum.290, %reduce_sum.291), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.407 (param_0.1384: f32[4,128], param_1.1570: f32[4,128]) -> (f32[], f32[]) { - %param_0.1384 = f32[4,128]{1,0:T(4,128)} parameter(0) - %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.261 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.413, %bitcast.413), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1216 = f32[]{:T(128)} constant(0) - %reduce.214 = f32[]{:T(128)} reduce(%square.261, %constant.1216), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1570 = f32[4,128]{1,0:T(4,128)} parameter(1) - %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.264.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.417.clone.1, %bitcast.417.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.215.clone.1 = f32[]{:T(128)} reduce(%square.264.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.407 (param_0.1383: f32[4,128], param_1.1563: f32[4,128]) -> (f32[], f32[]) { + %param_0.1383 = f32[4,128]{1,0:T(4,128)} parameter(0) + %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1965 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.413, %bitcast.413), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1214 = f32[]{:T(128)} constant(0) + %reduce.214 = f32[]{:T(128)} reduce(%mul.1965, %constant.1214), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1563 = f32[4,128]{1,0:T(4,128)} parameter(1) + %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %mul.1968.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.417.clone.1, %bitcast.417.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.215.clone.1 = f32[]{:T(128)} reduce(%mul.1968.clone.1, %constant.1214), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.170 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.214, %reduce.215.clone.1) } @@ -1204,31 +1204,31 @@ StackFrames ROOT %reduce_sum.400 = f32[]{:T(128)} add(%reduce_sum.395, %reduce_sum.396), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.410 (param_0.1391: bf16[4,128], param_1.1576: f32[4,128], param_2.1329: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { +%fused_computation.410 (param_0.1390: bf16[4,128], param_1.1569: f32[4,128], param_2.1326: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { %param_3.932 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) %constant.1170.clone.1 = s32[]{:T(128)} constant(0) %broadcast.901.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1170.clone.1), dimensions={}, metadata={op_name="broadcast.95"} %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.932, %broadcast.901.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1576 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1576), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1391 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1569), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_0.1390 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1390), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} %add.927 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %square.269 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %add.927), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} - %constant.1224 = f32[]{:T(128)} constant(0) - %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1224), dimensions={}, metadata={op_name="broadcast.99"} - %mul.1913 = f32[4,128]{1,0:T(4,128)} multiply(%square.269, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %mul.1893 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1913, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.216 = f32[]{:T(128)} reduce(%mul.1893, %constant.1224), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1329 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1329), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} - %add.904.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1913), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %mul.1894.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.904.clone.1, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1894.clone.1, %constant.1224), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %mul.1911.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %broadcast.831), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %square.193 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %add.927), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.1222 = f32[]{:T(128)} constant(0) + %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1222), dimensions={}, metadata={op_name="broadcast.99"} + %mul.1989 = f32[4,128]{1,0:T(4,128)} multiply(%square.193, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1969 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1989, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.216 = f32[]{:T(128)} reduce(%mul.1969, %constant.1222), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_2.1326 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1326), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.904.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1989), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1970.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.904.clone.1, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1970.clone.1, %constant.1222), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1987.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %broadcast.831), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %constant.1068.clone.1 = f32[]{:T(128)} constant(1) %add.922.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1068.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - %add.915.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1911.clone.1, %add.922.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.915.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1987.clone.1, %add.922.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}, pred[4,128]{1,0:T(4,128)(4,1)S(1)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%reduce.216, %reduce.219.clone.1, %ne.6.clone.1, %add.915.clone.1) } @@ -1244,39 +1244,39 @@ StackFrames ROOT %reduce_sum.379 = f32[]{:T(128)} add(%reduce_sum.374, %reduce_sum.375), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.411 (param_0.1370: f32[128,4], param_1.1558: f32[], param_2.1316: f32[], param_3.920: f32[], param_4.558: f32[128,4], param_5.470: f32[], param_6.360: f32[4,128], param_7.203: pred[], param_8.120: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1370 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.411 (param_0.1369: f32[128,4], param_1.1551: f32[], param_2.1313: f32[], param_3.920: f32[], param_4.556: f32[128,4], param_5.469: f32[], param_6.358: f32[4,128], param_7.198: pred[], param_8.115: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1369 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.920 = f32[]{:T(128)S(6)} parameter(3) - %mul.1943.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.920), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.203 = pred[]{:T(512)S(6)} parameter(7) - %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.360 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.470 = f32[]{:T(128)} parameter(5) - %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2019.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.920), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.198 = pred[]{:T(512)S(6)} parameter(7) + %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.198), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.358 = f32[4,128]{1,0:T(4,128)} parameter(6) + %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.469 = f32[]{:T(128)} parameter(5) + %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.875.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.468.clone.1, %div.876.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.275.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.276.clone.1, %bitcast.468.clone.1, %div.875.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1104.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.856.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1104.clone.1), dimensions={}, metadata={op_name="broadcast.78"} - %mul.1947.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.120 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.2023.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.115 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1108.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.855.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1108.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.1946.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.951.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1947.clone.1, %mul.1946.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) - %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2022.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.115, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.951.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2023.clone.1, %mul.2022.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1313 = f32[]{:T(128)S(6)} parameter(2) + %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1313), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.67.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %select_n.275.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1107.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.854.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1107.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1945.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.854.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.558 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.2021.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.854.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.556 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1106.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.853.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1106.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1944.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.558, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.950.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1945.clone.1, %mul.1944.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) - %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2020.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.556, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.950.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2021.clone.1, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1551 = f32[]{:T(128)S(6)} parameter(1) + %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1551), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.870.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.950.clone.1, %div.871.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.64.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.870.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1105.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1284,14 +1284,14 @@ StackFrames %add.949.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.851.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.428.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.872.clone.1, %add.949.clone.1), metadata={op_name="multiply.59"} %div.869.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.951.clone.1, %multiply.428.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1942.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1370, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.948.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.869.clone.1, %mul.1942.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1941.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1943.clone.1, %add.948.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1370, %mul.1941.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.265 = f32[128,4]{0,1:T(4,128)} multiply(%add.947.clone.1, %add.947.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1202 = f32[]{:T(128)} constant(0) - %reduce.217 = f32[]{:T(128)} reduce(%square.265, %constant.1202), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1202), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2018.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1369, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.948.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.869.clone.1, %mul.2018.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2017.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.2019.clone.1, %add.948.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1369, %mul.2017.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.189 = f32[128,4]{0,1:T(4,128)} multiply(%add.947.clone.1, %add.947.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1200 = f32[]{:T(128)} constant(0) + %reduce.217 = f32[]{:T(128)} reduce(%square.189, %constant.1200), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.159 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.217, %add.947.clone.1, %add.950.clone.1, %add.951.clone.1, %reduce.221.clone.1) } @@ -1307,39 +1307,39 @@ StackFrames ROOT %reduce_sum.361 = f32[]{:T(128)} add(%reduce_sum.359, %reduce_sum.360), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.412 (param_0.1373: f32[128,4], param_1.1561: f32[], param_2.1319: f32[], param_3.923: f32[], param_4.561: f32[128,4], param_5.473: f32[], param_6.363: f32[4,128], param_7.206: pred[], param_8.123: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1373 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.412 (param_0.1372: f32[128,4], param_1.1554: f32[], param_2.1316: f32[], param_3.923: f32[], param_4.559: f32[128,4], param_5.472: f32[], param_6.361: f32[4,128], param_7.201: pred[], param_8.118: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1372 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.923 = f32[]{:T(128)S(6)} parameter(3) - %mul.1970.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.923), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.206 = pred[]{:T(512)S(6)} parameter(7) - %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.363 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.473 = f32[]{:T(128)} parameter(5) - %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2046.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.923), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.201 = pred[]{:T(512)S(6)} parameter(7) + %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.201), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.361 = f32[4,128]{1,0:T(4,128)} parameter(6) + %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.472 = f32[]{:T(128)} parameter(5) + %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.899.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.474.clone.1, %div.900.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.287.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.288.clone.1, %bitcast.474.clone.1, %div.899.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1122.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.866.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1122.clone.1), dimensions={}, metadata={op_name="broadcast.78"} - %mul.1974.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.123 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.2050.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.118 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1126.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.865.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1126.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.1973.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.123, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.968.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1974.clone.1, %mul.1973.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) - %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2049.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.118, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.968.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2050.clone.1, %mul.2049.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) + %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.70.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %select_n.287.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1125.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.864.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1125.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1972.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.864.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.561 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.2048.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.864.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.559 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1124.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.863.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1124.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1971.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.967.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1972.clone.1, %mul.1971.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) - %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2047.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.559, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.967.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2048.clone.1, %mul.2047.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1554 = f32[]{:T(128)S(6)} parameter(1) + %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1554), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.894.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.967.clone.1, %div.895.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.67.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.894.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1123.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1347,35 +1347,35 @@ StackFrames %add.966.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.67.clone.1, %broadcast.861.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.431.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.896.clone.1, %add.966.clone.1), metadata={op_name="multiply.56"} %div.893.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.968.clone.1, %multiply.431.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1969.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1373, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.965.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.893.clone.1, %mul.1969.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1968.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1970.clone.1, %add.965.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1373, %mul.1968.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.266 = f32[128,4]{0,1:T(4,128)} multiply(%add.964.clone.1, %add.964.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1205 = f32[]{:T(128)} constant(0) - %reduce.218 = f32[]{:T(128)} reduce(%square.266, %constant.1205), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2045.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1372, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.965.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.893.clone.1, %mul.2045.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2044.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.2046.clone.1, %add.965.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1372, %mul.2044.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.190 = f32[128,4]{0,1:T(4,128)} multiply(%add.964.clone.1, %add.964.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1203 = f32[]{:T(128)} constant(0) + %reduce.218 = f32[]{:T(128)} reduce(%square.190, %constant.1203), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1203), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.160 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.218, %add.964.clone.1, %add.967.clone.1, %add.968.clone.1, %reduce.222.clone.1) } -%fused_computation.421 (param_0.1201: f32[4,128], param_1.1323: f32[4,128]) -> f32[4,128] { - %param_0.1201 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1323 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.421 (param_0.1200: f32[4,128], param_1.1320: f32[4,128]) -> f32[4,128] { + %param_0.1200 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1320 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.1045 = f32[]{:T(128)} constant(0.00048828125) %broadcast.837 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1045), dimensions={}, metadata={op_name="broadcast.399"} - %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1323, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1320, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1043 = f32[]{:T(128)} constant(1e-06) %add.935 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1043), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.934 = f32[4,128]{1,0:T(4,128)} add(%div.767, %add.935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %rsqrt.168 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.934), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} %div.754 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.168, %add.934), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1040 = f32[]{:T(128)} constant(-0.5) - %mul.1919 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1040), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1910 = f32[4,128]{1,0:T(4,128)} multiply(%div.754, %mul.1919), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1909 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1201, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1995 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1040), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1986 = f32[4,128]{1,0:T(4,128)} multiply(%div.754, %mul.1995), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1985 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1200, %mul.1986), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.1039 = f32[]{:T(128)} constant(0.0009765625) - %mul.1918 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1039), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - ROOT %mul.1908 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1909, %mul.1918), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1994 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1039), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + ROOT %mul.1984 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1985, %mul.1994), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} } %region_0.1 (reduce_sum.137: s32[], reduce_sum.138: s32[]) -> s32[] { @@ -1384,31 +1384,31 @@ StackFrames ROOT %reduce_sum.139 = s32[]{:T(128)} add(%reduce_sum.137, %reduce_sum.138), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.425 (param_0.1220: pred[4,128]) -> s32[] { - %param_0.1220 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1220), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} +%fused_computation.425 (param_0.1219: pred[4,128]) -> s32[] { + %param_0.1219 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1219), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} %constant.1066 = s32[]{:T(128)} constant(0) ROOT %reduce.220 = s32[]{:T(128)} reduce(%convert_element_type.1403, %constant.1066), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.428 (param_0.1203: f32[4,128]) -> f32[4,128] { - %param_0.1203 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.428 (param_0.1202: f32[4,128]) -> f32[4,128] { + %param_0.1202 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1046 = f32[]{:T(128)} constant(0.00048828125) %broadcast.829 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1046), dimensions={}, metadata={op_name="broadcast.399"} - %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1203, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1202, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1044 = f32[]{:T(128)} constant(1e-06) %add.924 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1044), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.921 = f32[4,128]{1,0:T(4,128)} add(%div.759, %add.924), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} ROOT %rsqrt.166 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.921), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} } -%fused_computation.429 (param_0.1204: pred[4,128], param_1.1575: f32[]) -> f32[4,128] { - %param_0.1204 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1575 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1575), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} - %constant.1223 = f32[]{:T(128)} constant(0) - %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.99"} - ROOT %mul.1920 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1204, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.429 (param_0.1203: pred[4,128], param_1.1568: f32[]) -> f32[4,128] { + %param_0.1203 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) + %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} + %constant.1221 = f32[]{:T(128)} constant(0) + %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1221), dimensions={}, metadata={op_name="broadcast.99"} + ROOT %mul.1996 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1203, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} } %fused_computation.431 () -> f32[64] { @@ -1417,43 +1417,43 @@ StackFrames %iota.46 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/layers/iota" stack_frame_id=0} %constant.1048 = s32[]{:T(128)} constant(2) %broadcast.839 = s32[64]{0:T(128)} broadcast(%constant.1048), dimensions={}, metadata={op_name="broadcast.391"} - %mul.1921 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.839), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} - %convert_element_type.1404 = f32[64]{0:T(128)} convert(%mul.1921), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %mul.1997 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.839), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} + %convert_element_type.1404 = f32[64]{0:T(128)} convert(%mul.1997), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %constant.1047 = f32[]{:T(128)} constant(0.0078125) %broadcast.838 = f32[64]{0:T(128)} broadcast(%constant.1047), dimensions={}, metadata={op_name="broadcast.392"} %div.768 = f32[64]{0:T(128)} multiply(%convert_element_type.1404, %broadcast.838), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.840, %div.768), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} } -%fused_computation.432 (param_0.1218: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.1218 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1218), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} +%fused_computation.432 (param_0.1217: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { + %param_0.1217 = s32[4,128]{1,0:T(4,128)} parameter(0) + %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1217), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %bitcast.418 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} ROOT %tuple.162 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.418, %convert_element_type.1405) } -%fused_computation.435 (param_0.1360: f32[2048,4]) -> bf16[4,2048] { - %param_0.1360 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.531 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.531) +%fused_computation.435 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { + %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) + %bitcast.533 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.533) } -%fused_computation.436 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { - %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.530 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.530) +%fused_computation.436 (param_0.1358: f32[2048,4]) -> bf16[4,2048] { + %param_0.1358 = f32[2048,4]{0,1:T(4,128)} parameter(0) + %bitcast.532 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.532) } -%fused_computation.437 (param_0.1361: f32[128,4]) -> bf16[4,128] { - %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.532 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.532) +%fused_computation.437 (param_0.1360: f32[128,4]) -> bf16[4,128] { + %param_0.1360 = f32[128,4]{0,1:T(4,128)} parameter(0) + %bitcast.534 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.534) } -%fused_computation.438 (param_0.1362: f32[128,4]) -> bf16[4,128] { - %param_0.1362 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.533 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1362), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.533) +%fused_computation.438 (param_0.1361: f32[128,4]) -> bf16[4,128] { + %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) + %bitcast.535 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.535) } %region_8.11 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { @@ -1462,40 +1462,40 @@ StackFrames ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.287.clone.clone (param_0.1346: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1346 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.526 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1346), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%fused_computation.287.clone.clone (param_0.1345: bf16[151936,2048]) -> bf16[151936,2048,1] { + %param_0.1345 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1345), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} } -%fused_computation.368.clone.clone (param_0.1347: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1281: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1281 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.476 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1281), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1438 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1347 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2067 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1347), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2066 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1438, %mul.2067), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2066), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.476, %convert_element_type.1437), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.368.clone.clone (param_0.1346: f32[4,128], param_1.1535: bf16[4,128,2048], param_2.1278: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1278 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1278), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1535 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1432 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1535), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1346 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2135 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1346), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.2134 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1432, %mul.2135), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1431 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2134), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.474 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.475, %convert_element_type.1431), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.439 (param_0.1363: bf16[151936,2048], param_1.1551: f32[4,128], param_2.1305: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { - %param_1.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1305 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) +%fused_computation.439 (param_0.1362: bf16[151936,2048], param_1.1544: f32[4,128], param_2.1302: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { + %param_1.1544 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1302 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) %param_3.913 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.270.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1551, %param_2.1305, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1363 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %fusion.253.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1363), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.270.clone.1, %fusion.253.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %constant.1195 = bf16[]{:T(256)} constant(-inf) - %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1195), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1544, %param_2.1302, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1362 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %fusion.252.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1362), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.269.clone.1, %fusion.252.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %constant.1193 = bf16[]{:T(256)} constant(-inf) + %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1193), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} ROOT %tuple.164 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,151936]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.223, %convolution.85.clone.1) } -%fused_computation.440 (param_0.1358: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.1358 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %bitcast.529 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.529) +%fused_computation.440 (param_0.1357: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { + %param_0.1357 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %bitcast.531 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1357), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.531) } %convert_element_type.767.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { @@ -1504,13 +1504,13 @@ StackFrames ROOT %add.755 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.155.clone.clone (param_0.1534: bf16[4,2048], param_1.1687: s32[]) -> bf16[2048] { - %param_0.1534 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) - %constant.1361 = s32[]{:T(128)} constant(0) - %dynamic_slice.388 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1687, %constant.1361), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1362 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.388, %constant.1362), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.155.clone.clone (param_0.1533: bf16[4,2048], param_1.1680: s32[]) -> bf16[2048] { + %param_0.1533 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1680 = s32[]{:T(128)S(6)} parameter(1) + %constant.1359 = s32[]{:T(128)} constant(0) + %dynamic_slice.384 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1533, %param_1.1680, %constant.1359), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1360 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.384, %constant.1360), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %region_14.16 (reduce_sum.204: f32[], reduce_sum.205: f32[]) -> f32[] { @@ -1519,25 +1519,25 @@ StackFrames ROOT %reduce_sum.206 = f32[]{:T(128)} add(%reduce_sum.204, %reduce_sum.205), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.58.clone.clone (param_0.1535: bf16[4,4,128,2048], param_1.1688: s32[]) -> f32[4,128] { - %param_0.1535 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1688 = s32[]{:T(128)S(6)} parameter(1) - %constant.1363 = s32[]{:T(128)} constant(0) - %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1535, %param_1.1688, %constant.1363, %constant.1363, %constant.1363), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1564 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.633), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.280 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1564, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1364 = f32[]{:T(128)} constant(0) - ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.280, %constant.1364), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} -} - -%fused_computation.179.clone.1.clone (param_0.1536: f32[4,128]) -> f32[4,128] { - %param_0.1536 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1366 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1366), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1536, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1365 = f32[]{:T(128)} constant(1e-06) - %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1365), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} +%fused_computation.58.clone.clone (param_0.1534: bf16[4,4,128,2048], param_1.1681: s32[]) -> f32[4,128] { + %param_0.1534 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1681 = s32[]{:T(128)S(6)} parameter(1) + %constant.1361 = s32[]{:T(128)} constant(0) + %dynamic_slice.385 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1681, %constant.1361, %constant.1361, %constant.1361), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.635 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1558 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.635), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.204 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1558, %convert_element_type.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1362 = f32[]{:T(128)} constant(0) + ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.204, %constant.1362), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} +} + +%fused_computation.179.clone.1.clone (param_0.1535: f32[4,128]) -> f32[4,128] { + %param_0.1535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %constant.1364 = f32[]{:T(128)} constant(0.00048828125) + %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1364), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1535, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1363 = f32[]{:T(128)} constant(1e-06) + %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1363), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.1039 = f32[4,128]{1,0:T(4,128)} add(%div.999, %closed_call.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.181 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1039), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } @@ -1548,158 +1548,158 @@ StackFrames ROOT %reduce_sum.212 = f32[]{:T(128)} add(%reduce_sum.207, %reduce_sum.211), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1550: bf16[4,2048,16,128], param_1.1698: s32[]) -> bf16[2048,16,128,1] { - %param_0.1550 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1698 = s32[]{:T(128)S(6)} parameter(1) - %constant.1377 = s32[]{:T(128)} constant(0) - %dynamic_slice.395 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1550, %param_1.1698, %constant.1377, %constant.1377, %constant.1377), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.644 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.114.clone.clone.clone.clone (param_0.1551: f32[4,128], param_1.1699: bf16[4,4,128,2048], param_2.1405: s32[], param_3.982: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.982 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.571 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1699 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1405 = s32[]{:T(128)S(6)} parameter(2) - %constant.1378 = s32[]{:T(128)} constant(0) - %dynamic_slice.396 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1699, %param_2.1405, %constant.1378, %constant.1378, %constant.1378), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.646 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.646), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2256 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1551), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2255 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %mul.2256), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2255), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.571, %convert_element_type.1574), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.645 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.61.clone.clone (param_0.1552: bf16[4,2048,16,128], param_1.1700: s32[], param_2.1406: f32[4,128], param_3.983: bf16[4,4,128,2048], param_4.604: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { - %param_2.1406 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.983 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1700 = s32[]{:T(128)S(6)} parameter(1) - %param_4.604 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1406, %param_3.983, %param_1.1700, %param_4.604), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1552 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1552, %param_1.1700), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.44.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1576 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.44.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.282 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1576, %convert_element_type.1576), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1379 = f32[]{:T(128)} constant(0) - %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.282, %constant.1379), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.44.clone.3) -} - -%fused_computation.151.clone.1.clone (param_0.1553: f32[4,128,16]) -> f32[4,128,16] { - %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1380 = f32[]{:T(128)} constant(0.0078125) - %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1380), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1553, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1381 = f32[]{:T(128)} constant(1e-06) - %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} +%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1549: bf16[4,2048,16,128], param_1.1691: s32[]) -> bf16[2048,16,128,1] { + %param_0.1549 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1691 = s32[]{:T(128)S(6)} parameter(1) + %constant.1375 = s32[]{:T(128)} constant(0) + %dynamic_slice.391 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1691, %constant.1375, %constant.1375, %constant.1375), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.646 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.114.clone.clone.clone.clone (param_0.1550: f32[4,128], param_1.1692: bf16[4,4,128,2048], param_2.1403: s32[], param_3.983: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.983 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.983), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1692 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1403 = s32[]{:T(128)S(6)} parameter(2) + %constant.1376 = s32[]{:T(128)} constant(0) + %dynamic_slice.392 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1692, %param_2.1403, %constant.1376, %constant.1376, %constant.1376), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.648 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1569 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.648), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1550 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2324 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2323 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1569, %mul.2324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1568 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.569 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.570, %convert_element_type.1568), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.647 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.61.clone.clone (param_0.1551: bf16[4,2048,16,128], param_1.1693: s32[], param_2.1404: f32[4,128], param_3.984: bf16[4,4,128,2048], param_4.603: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { + %param_2.1404 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.984 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1693 = s32[]{:T(128)S(6)} parameter(1) + %param_4.603 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1404, %param_3.984, %param_1.1693, %param_4.603), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1551 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1551, %param_1.1693), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.46.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %convert_element_type.1570 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.46.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.206 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1570, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1377 = f32[]{:T(128)} constant(0) + %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.206, %constant.1377), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.46.clone.3) +} + +%fused_computation.151.clone.1.clone (param_0.1552: f32[4,128,16]) -> f32[4,128,16] { + %param_0.1552 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) + %constant.1378 = f32[]{:T(128)} constant(0.0078125) + %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1378), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1552, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1379 = f32[]{:T(128)} constant(1e-06) + %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1379), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} %add.1043 = f32[4,128,16]{1,2,0:T(8,128)} add(%div.1001, %add.1044), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.183 = f32[4,128,16]{1,2,0:T(8,128)S(1)} rsqrt(%add.1043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.182.clone.clone (param_0.1549: bf16[4,128], param_1.1697: s32[]) -> bf16[128] { - %param_0.1549 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) - %constant.1376 = s32[]{:T(128)} constant(0) - %dynamic_slice.394 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1697, %constant.1376), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.643 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.394), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.182.clone.clone (param_0.1548: bf16[4,128], param_1.1690: s32[]) -> bf16[128] { + %param_0.1548 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1690 = s32[]{:T(128)S(6)} parameter(1) + %constant.1374 = s32[]{:T(128)} constant(0) + %dynamic_slice.390 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1548, %param_1.1690, %constant.1374), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.645 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.121.clone.1.clone (param_0.1554: f32[4,128,16], param_1.1701: bf16[4,128,16,128], param_2.1407: bf16[128]) -> bf16[4,128,16,128] { - %param_2.1407 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1407), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1701 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1578 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1554 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2258 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1554), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2257 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1578, %mul.2258), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1577 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2257), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.573, %convert_element_type.1577), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.121.clone.1.clone (param_0.1553: f32[4,128,16], param_1.1694: bf16[4,128,16,128], param_2.1405: bf16[128]) -> bf16[4,128,16,128] { + %param_2.1405 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1405), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1694 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1572 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1694), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) + %mul.2326 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1553), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2325 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1572, %mul.2326), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1571 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.571 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.572, %convert_element_type.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.90.clone.clone (param_0.1555: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { - %param_0.1555 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.90.clone.clone (param_0.1554: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { + %param_0.1554 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1554), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.129 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1554), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.209 = (bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) } %fused_computation.187.clone.clone () -> f32[64] { - %constant.1355 = f32[]{:T(128)} constant(1e+06) - %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1355), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %constant.1353 = f32[]{:T(128)} constant(1e+06) + %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1353), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} - %constant.1354 = s32[]{:T(128)} constant(2) - %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2242 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1562 = f32[64]{0:T(128)} convert(%mul.2242), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %constant.1356 = f32[]{:T(128)} constant(0.0078125) - %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1356), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1562, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1352 = s32[]{:T(128)} constant(2) + %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1352), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %mul.2310 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1556 = f32[64]{0:T(128)} convert(%mul.2310), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %constant.1354 = f32[]{:T(128)} constant(0.0078125) + %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1556, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.104, %div.995), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} } -%fused_computation.143.clone.clone (param_0.1529: f32[64], param_1.1683: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1683 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1683), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1529 = f32[64]{0:T(128)S(1)} parameter(0) - %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1529), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} +%fused_computation.143.clone.clone (param_0.1528: f32[64], param_1.1676: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1676 = f32[4,128]{1,0:T(4,128)} parameter(1) + %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1676), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %param_0.1528 = f32[64]{0:T(128)S(1)} parameter(0) + %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1528), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %div.996 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.998, %div.997), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1563 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1557 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} %convert_element_type.1189.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1563, %convert_element_type.1189.clone.3) + ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1557, %convert_element_type.1189.clone.3) } -%fused_computation.146.clone.1.clone (param_0.1530: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1530 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1357 = bf16[]{:T(256)} constant(-inf) - %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} +%fused_computation.146.clone.1.clone (param_0.1529: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1529 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) + %constant.1355 = bf16[]{:T(256)} constant(-inf) + %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1529, %constant.1355), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1529, %constant.1355), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.53 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.69, %pad.68), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.630 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.632 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.145.clone.1.clone (param_0.1545: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1545 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1374 = bf16[]{:T(256)} constant(-inf) - %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} +%fused_computation.145.clone.1.clone (param_0.1544: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1544 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) + %constant.1372 = bf16[]{:T(256)} constant(-inf) + %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1544, %constant.1372), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1544, %constant.1372), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.54 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.71, %pad.70), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.641 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.643 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.94.clone.clone (param_0.1556: bf16[4,128,16,64], param_1.1702: bf16[4,128,16,64], param_2.1408: bf16[4,128,128], param_3.984: bf16[4,128,128], param_4.605: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.384: bf16[128]) -> bf16[4,16,128,128] { - %param_6.384 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.575 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.384), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.94.clone.clone (param_0.1555: bf16[4,128,16,64], param_1.1695: bf16[4,128,16,64], param_2.1406: bf16[4,128,128], param_3.985: bf16[4,128,128], param_4.604: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.383: bf16[128]) -> bf16[4,16,128,128] { + %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} %param_5.499 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1580 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.605 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2265 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.605), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2264 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1580, %mul.2265), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1579 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.575, %convert_element_type.1579), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.984 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2263 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.984), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2261 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %mul.2263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1702 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1382 = bf16[]{:T(256)} constant(-inf) - %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1702, %constant.1382), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1556 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1556, %constant.1382), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %convert_element_type.1574 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_4.604 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) + %mul.2333 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.604), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2332 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1574, %mul.2333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2332), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %convert_element_type.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.985 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2331 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.985), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2329 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.573, %mul.2331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1695 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1380 = bf16[]{:T(256)} constant(-inf) + %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1695, %constant.1380), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1555 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1555, %constant.1380), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.56 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.75, %pad.74), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1408 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2262 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1408), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2260 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2262), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2261, %mul.2260), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %constant.1383 = bf16[]{:T(256)} constant(0.08838) - %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1383), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2259 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - ROOT %bitcast.647 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2259), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_2.1406 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2330 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1406), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2328 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2329, %mul.2328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + %constant.1381 = bf16[]{:T(256)} constant(0.08838) + %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %mul.2327 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.649 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } %region_16.18 (reduce_sum.213: f32[], reduce_sum.214: f32[]) -> f32[] { @@ -1708,159 +1708,159 @@ StackFrames ROOT %reduce_sum.218 = f32[]{:T(128)} add(%reduce_sum.213, %reduce_sum.214), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1541: bf16[4,2048,8,128], param_1.1692: s32[]) -> bf16[2048,8,128,1] { - %param_0.1541 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1692 = s32[]{:T(128)S(6)} parameter(1) - %constant.1369 = s32[]{:T(128)} constant(0) - %dynamic_slice.392 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1541, %param_1.1692, %constant.1369, %constant.1369, %constant.1369), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.638 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.113.clone.clone.clone.clone (param_0.1542: f32[4,128], param_1.1693: bf16[4,4,128,2048], param_2.1401: s32[], param_3.979: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.979 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.979), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1693 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1401 = s32[]{:T(128)S(6)} parameter(2) - %constant.1370 = s32[]{:T(128)} constant(0) - %dynamic_slice.393 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1693, %param_2.1401, %constant.1370, %constant.1370, %constant.1370), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.640 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1568 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.640), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1542 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2246 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1542), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2245 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1568, %mul.2246), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.565, %convert_element_type.1567), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.639 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.84.clone.clone (param_0.1543: bf16[4,2048,8,128], param_1.1694: s32[], param_2.1402: f32[4,128], param_3.980: bf16[4,4,128,2048], param_4.602: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { - %param_2.1402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.980 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1694 = s32[]{:T(128)S(6)} parameter(1) - %param_4.602 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1402, %param_3.980, %param_1.1694, %param_4.602), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1543 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1543, %param_1.1694), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1540: bf16[4,2048,8,128], param_1.1685: s32[]) -> bf16[2048,8,128,1] { + %param_0.1540 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) + %constant.1367 = s32[]{:T(128)} constant(0) + %dynamic_slice.388 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1540, %param_1.1685, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.640 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.113.clone.clone.clone.clone (param_0.1541: f32[4,128], param_1.1686: bf16[4,4,128,2048], param_2.1399: s32[], param_3.980: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.980 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.980), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1686 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) + %constant.1368 = s32[]{:T(128)} constant(0) + %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1686, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.642 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1562 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.642), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1541 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2314 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1541), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2313 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1562, %mul.2314), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1561 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2313), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.564, %convert_element_type.1561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.641 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.84.clone.clone (param_0.1542: bf16[4,2048,8,128], param_1.1687: s32[], param_2.1400: f32[4,128], param_3.981: bf16[4,4,128,2048], param_4.601: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { + %param_2.1400 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.981 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) + %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1400, %param_3.981, %param_1.1687, %param_4.601), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1542 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1542, %param_1.1687), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %convolution.50.clone.3 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.73.clone.3, %fusion.87.clone.3), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1569 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.281 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1569, %convert_element_type.1569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1371 = f32[]{:T(128)} constant(0) - %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.281, %constant.1371), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %convert_element_type.1563 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.205 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1563, %convert_element_type.1563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1369 = f32[]{:T(128)} constant(0) + %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.205, %constant.1369), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.206 = (f32[4,128,8]{1,2,0:T(8,128)S(1)}, bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.246, %convolution.50.clone.3) } -%fused_computation.154.clone.1.clone (param_0.1544: f32[4,128,8]) -> f32[4,128,8] { - %param_0.1544 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1372 = f32[]{:T(128)} constant(0.0078125) - %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1372), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1544, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1373 = f32[]{:T(128)} constant(1e-06) - %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1373), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} +%fused_computation.154.clone.1.clone (param_0.1543: f32[4,128,8]) -> f32[4,128,8] { + %param_0.1543 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) + %constant.1370 = f32[]{:T(128)} constant(0.0078125) + %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1370), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1543, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1371 = f32[]{:T(128)} constant(1e-06) + %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1371), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} %add.1040 = f32[4,128,8]{1,2,0:T(8,128)} add(%div.1000, %add.1041), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.182 = f32[4,128,8]{1,2,0:T(8,128)S(1)} rsqrt(%add.1040), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.184.clone.clone (param_0.1528: bf16[4,128], param_1.1682: s32[]) -> bf16[128] { - %param_0.1528 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) - %constant.1353 = s32[]{:T(128)} constant(0) - %dynamic_slice.385 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1528, %param_1.1682, %constant.1353), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.629 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.139.clone.1.clone (param_0.1546: f32[4,128,8], param_1.1695: bf16[4,128,8,128], param_2.1403: bf16[128]) -> bf16[4,128,8,128] { - %param_2.1403 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1403), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1695 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1571 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1695), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1546 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2248 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1546), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2247 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1571, %mul.2248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1570 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2247), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.567, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.126.clone.clone (param_0.1547: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1547 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - ROOT %tuple.207 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) +%fused_computation.184.clone.clone (param_0.1527: bf16[4,128], param_1.1675: s32[]) -> bf16[128] { + %param_0.1527 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1675 = s32[]{:T(128)S(6)} parameter(1) + %constant.1351 = s32[]{:T(128)} constant(0) + %dynamic_slice.381 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1527, %param_1.1675, %constant.1351), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.631 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.129.clone.clone (param_0.1548: bf16[4,128,8,64], param_1.1696: bf16[4,128,8,64], param_2.1404: bf16[4,128,128], param_3.981: bf16[4,128,128], param_4.603: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.383: bf16[128]) -> bf16[4,8,128,128] { - %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.569 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_5.498 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1573 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.603 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2254 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.603), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2253 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1573, %mul.2254), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1572 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2253), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.569, %convert_element_type.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.981 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2252 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.981), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2250 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %mul.2252), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1696 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1375 = bf16[]{:T(256)} constant(-inf) - %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1696, %constant.1375), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1548 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1548, %constant.1375), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.55 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.73, %pad.72), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1404 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2251 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1404), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2249 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2250, %mul.2249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.642 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.169.clone.clone (param_0.1537: bf16[4,2048,8,128], param_1.1689: s32[]) -> bf16[1,2048,8,128] { - %param_0.1537 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1689 = s32[]{:T(128)S(6)} parameter(1) - %constant.1367 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.390 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1537, %param_1.1689, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +%fused_computation.139.clone.1.clone (param_0.1545: f32[4,128,8], param_1.1688: bf16[4,128,8,128], param_2.1401: bf16[128]) -> bf16[4,128,8,128] { + %param_2.1401 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1401), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1688 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1565 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1688), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1545 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) + %mul.2316 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1545), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2315 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1565, %mul.2316), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1564 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2315), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.565 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.566, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1538: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { - %param_0.1538 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.204 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1538), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.634 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.204), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.126.clone.clone (param_0.1546: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { + %param_0.1546 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1546), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} + %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1546), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + ROOT %tuple.207 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) } -%fused_computation.111.clone.clone.clone.clone (param_0.1539: f32[4,128], param_1.1690: bf16[4,4,128,2048], param_2.1399: s32[], param_3.977: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.977 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.977), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1690 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) - %constant.1368 = s32[]{:T(128)} constant(0) - %dynamic_slice.391 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1690, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.636 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1566 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.636), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1539 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2244 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2243 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1566, %mul.2244), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2243), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.563, %convert_element_type.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.635 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.562), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.140.clone.clone (param_0.1540: bf16[1,2048,8,128], param_1.1691: f32[4,128], param_2.1400: bf16[4,4,128,2048], param_3.978: s32[], param_4.601: bf16[2048]) -> bf16[4,8,128,128] { - %param_1.1691 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1400 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.978 = s32[]{:T(128)S(6)} parameter(3) - %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.373 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1691, %param_2.1400, %param_3.978, %param_4.601), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1540 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.372 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1540), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.373, %fusion.372), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.637 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.129.clone.clone (param_0.1547: bf16[4,128,8,64], param_1.1689: bf16[4,128,8,64], param_2.1402: bf16[4,128,128], param_3.982: bf16[4,128,128], param_4.602: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.382: bf16[128]) -> bf16[4,8,128,128] { + %param_6.382 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.382), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_5.498 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) + %convert_element_type.1567 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_4.602 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) + %mul.2322 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.602), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2321 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1567, %mul.2322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %convert_element_type.1566), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.982 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2320 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2318 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.567, %mul.2320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1689 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1373 = bf16[]{:T(256)} constant(-inf) + %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1689, %constant.1373), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1547 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1547, %constant.1373), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %maximum.55 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.73, %pad.72), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_2.1402 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2319 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1402), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2317 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2319), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2318, %mul.2317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.644 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.188.clone.clone (param_0.1578: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { - %param_0.1578 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1578), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.660 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.660, %slice.11) +%fused_computation.169.clone.clone (param_0.1536: bf16[4,2048,8,128], param_1.1682: s32[]) -> bf16[1,2048,8,128] { + %param_0.1536 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) + %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) + %constant.1365 = s32[]{:T(128)} constant(0) + ROOT %dynamic_slice.386 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1536, %param_1.1682, %constant.1365, %constant.1365, %constant.1365), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1537: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { + %param_0.1537 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %copy.208 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1537), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} + ROOT %bitcast.636 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.208), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.111.clone.clone.clone.clone (param_0.1538: f32[4,128], param_1.1683: bf16[4,4,128,2048], param_2.1397: s32[], param_3.978: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.978 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.978), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1683 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1397 = s32[]{:T(128)S(6)} parameter(2) + %constant.1366 = s32[]{:T(128)} constant(0) + %dynamic_slice.387 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1683, %param_2.1397, %constant.1366, %constant.1366, %constant.1366), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.638 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1560 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.638), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1538 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2312 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1538), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2311 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1560, %mul.2312), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1559 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2311), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.561 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.562, %convert_element_type.1559), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.637 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.140.clone.clone (param_0.1539: bf16[1,2048,8,128], param_1.1684: f32[4,128], param_2.1398: bf16[4,4,128,2048], param_3.979: s32[], param_4.600: bf16[2048]) -> bf16[4,8,128,128] { + %param_1.1684 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1398 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.979 = s32[]{:T(128)S(6)} parameter(3) + %param_4.600 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.372 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1684, %param_2.1398, %param_3.979, %param_4.600), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1539 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.371 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1539), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.372, %fusion.371), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + ROOT %bitcast.639 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.188.clone.clone (param_0.1577: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { + %param_0.1577 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) + %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1577), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.662 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} + ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.662, %slice.11) } %region_17.20 (reduce_sum.219: f32[], reduce_sum.220: f32[]) -> f32[] { @@ -1869,36 +1869,36 @@ StackFrames ROOT %reduce_sum.221 = f32[]{:T(128)} add(%reduce_sum.219, %reduce_sum.220), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,2048], param_1.1703: s32[]) -> bf16[16,128,2048,1] { - %param_0.1557 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) - %constant.1384 = s32[]{:T(128)} constant(0) - %dynamic_slice.397 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1557, %param_1.1703, %constant.1384, %constant.1384, %constant.1384), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.648 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.397), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1556: bf16[4,16,128,2048], param_1.1696: s32[]) -> bf16[16,128,2048,1] { + %param_0.1556 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1696 = s32[]{:T(128)S(6)} parameter(1) + %constant.1382 = s32[]{:T(128)} constant(0) + %dynamic_slice.393 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1556, %param_1.1696, %constant.1382, %constant.1382, %constant.1382), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.650 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1558: bf16[4,16,128,128]) -> bf16[4,128,16,128] { - %param_0.1558 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.649 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,128]) -> bf16[4,128,16,128] { + %param_0.1557 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.651 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.64.clone.clone (param_0.1559: bf16[4,16,128,2048], param_1.1704: s32[], param_2.1409: bf16[4,16,128,128], param_3.985: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { - %param_3.985 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1704 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.64.clone.clone (param_0.1558: bf16[4,16,128,2048], param_1.1697: s32[], param_2.1407: bf16[4,16,128,128], param_3.986: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { + %param_3.986 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) %constant.436.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.242.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.985, %param_1.1704, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.242.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1409 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1409), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1559 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1559, %param_1.1704), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.240.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.986, %param_1.1697, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.240.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %param_2.1407 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1407), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_0.1558 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1558, %param_1.1697), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %convolution.62.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.96.clone.3, %fusion.95.clone.3), window={size=1x16}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %bitcast.203.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %add.768.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.227.clone.3, %bitcast.203.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1581 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.283 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1581, %convert_element_type.1581), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1385 = f32[]{:T(128)} constant(0) - %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.283, %constant.1385), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.207 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %convert_element_type.1575), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1383 = f32[]{:T(128)} constant(0) + %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.207, %constant.1383), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.210 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.248, %add.768.clone.3) } @@ -1908,93 +1908,93 @@ StackFrames ROOT %add.754 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.156.clone.clone (param_0.1531: bf16[4,2048], param_1.1684: s32[]) -> bf16[2048] { - %param_0.1531 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1684 = s32[]{:T(128)S(6)} parameter(1) - %constant.1358 = s32[]{:T(128)} constant(0) - %dynamic_slice.386 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1684, %constant.1358), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1359 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.386, %constant.1359), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.156.clone.clone (param_0.1530: bf16[4,2048], param_1.1677: s32[]) -> bf16[2048] { + %param_0.1530 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1677 = s32[]{:T(128)S(6)} parameter(1) + %constant.1356 = s32[]{:T(128)} constant(0) + %dynamic_slice.382 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1530, %param_1.1677, %constant.1356), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1357 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.382, %constant.1357), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.13.clone.clone.clone (param_0.1532: bf16[4,6144,2048], param_1.1685: s32[]) -> bf16[6144,2048,1] { - %param_0.1532 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) - %constant.1360 = s32[]{:T(128)} constant(0) - %dynamic_slice.387 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1532, %param_1.1685, %constant.1360, %constant.1360), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.632 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.13.clone.clone.clone (param_0.1531: bf16[4,6144,2048], param_1.1678: s32[]) -> bf16[6144,2048,1] { + %param_0.1531 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1678 = s32[]{:T(128)S(6)} parameter(1) + %constant.1358 = s32[]{:T(128)} constant(0) + %dynamic_slice.383 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1678, %constant.1358, %constant.1358), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.634 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %bitcast_fusion.1.clone.clone (bitcast_input.4: bf16[4,128,2048]) -> bf16[4,128,2048] { - %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.631 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) -} - -%fused_computation.14.clone.clone (param_0.1533: bf16[4,128,2048], param_1.1686: bf16[4,6144,2048], param_2.1398: s32[]) -> bf16[6144,4,128] { - %param_1.1686 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1398 = s32[]{:T(128)S(6)} parameter(2) - %fusion.370 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1686, %param_2.1398), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1533 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %fusion.371 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1533), kind=kLoop, calls=%bitcast_fusion.1.clone.clone - ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.370, %fusion.371), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.180.clone.1.clone (param_0.1560: f32[4,128]) -> f32[4,128] { - %param_0.1560 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1387 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1387), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1560, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1386 = f32[]{:T(128)} constant(1e-06) - %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1386), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) +} + +%fused_computation.14.clone.clone (param_0.1532: bf16[4,128,2048], param_1.1679: bf16[4,6144,2048], param_2.1396: s32[]) -> bf16[6144,4,128] { + %param_1.1679 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1396 = s32[]{:T(128)S(6)} parameter(2) + %fusion.369 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1679, %param_2.1396), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1532 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.370 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1532), kind=kLoop, calls=%bitcast_fusion.1.clone.clone + ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.369, %fusion.370), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.180.clone.1.clone (param_0.1559: f32[4,128]) -> f32[4,128] { + %param_0.1559 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %constant.1385 = f32[]{:T(128)} constant(0.00048828125) + %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1385), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1559, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1384 = f32[]{:T(128)} constant(1e-06) + %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1384), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.1046 = f32[4,128]{1,0:T(4,128)} add(%div.1002, %closed_call.110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.184 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1046), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.12.clone.1.clone.clone (param_0.1564: bf16[4,2048,6144], param_1.1708: s32[]) -> bf16[2048,6144,1] { - %param_0.1564 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1708 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.12.clone.1.clone.clone (param_0.1563: bf16[4,2048,6144], param_1.1701: s32[]) -> bf16[2048,6144,1] { + %param_0.1563 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1701 = s32[]{:T(128)S(6)} parameter(1) + %constant.1387 = s32[]{:T(128)} constant(0) + %dynamic_slice.395 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1563, %param_1.1701, %constant.1387, %constant.1387), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.119.clone.3.clone.clone (param_0.1564: f32[4,128], param_1.1702: bf16[4,128,2048], param_2.1410: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1410 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1410), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1702 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1579 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1564 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2337 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1564), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2336 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1579, %mul.2337), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2336), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.577 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.578, %convert_element_type.1578), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.21.clone.clone (param_0.1565: bf16[4,2048,6144], param_1.1703: s32[], param_2.1411: f32[4,128], param_3.988: bf16[4,128,2048], param_4.606: bf16[2048]) -> bf16[4,128,6144] { + %param_2.1411 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.988 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.606 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.376 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1411, %param_3.988, %param_4.606), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1565 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) + %fusion.375 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1565, %param_1.1703), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.376, %fusion.375), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.11.clone.1.clone.clone (param_0.1567: bf16[4,2048,6144], param_1.1705: s32[]) -> bf16[2048,6144,1] { + %param_0.1567 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1705 = s32[]{:T(128)S(6)} parameter(1) %constant.1389 = s32[]{:T(128)} constant(0) - %dynamic_slice.399 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1564, %param_1.1708, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.651 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.119.clone.3.clone.clone (param_0.1565: f32[4,128], param_1.1709: bf16[4,128,2048], param_2.1412: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1412 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.579 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1412), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1709 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1585 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1565 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2269 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1565), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2268 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1585, %mul.2269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1584 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.579, %convert_element_type.1584), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.21.clone.clone (param_0.1566: bf16[4,2048,6144], param_1.1710: s32[], param_2.1413: f32[4,128], param_3.987: bf16[4,128,2048], param_4.607: bf16[2048]) -> bf16[4,128,6144] { - %param_2.1413 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.987 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.607 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.377 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1413, %param_3.987, %param_4.607), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1566 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1710 = s32[]{:T(128)S(6)} parameter(1) - %fusion.376 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1566, %param_1.1710), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.377, %fusion.376), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.11.clone.1.clone.clone (param_0.1568: bf16[4,2048,6144], param_1.1712: s32[]) -> bf16[2048,6144,1] { - %param_0.1568 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1712 = s32[]{:T(128)S(6)} parameter(1) - %constant.1391 = s32[]{:T(128)} constant(0) - %dynamic_slice.400 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1568, %param_1.1712, %constant.1391, %constant.1391), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.400), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.47.clone.1.clone.clone (param_0.1567: bf16[6144,4,128], param_1.1711: bf16[4,128,6144]) -> bf16[4,128,6144] { - %param_1.1711 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1390 = bf16[]{:T(256)} constant(1) - %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1390), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1711), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %dynamic_slice.396 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1567, %param_1.1705, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.655 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.47.clone.1.clone.clone (param_0.1566: bf16[6144,4,128], param_1.1704: bf16[4,128,6144]) -> bf16[4,128,6144] { + %param_1.1704 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1388 = bf16[]{:T(256)} constant(1) + %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1388), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} + %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.69 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.1047 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.1003 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.1047), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.2271 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1711, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %mul.2339 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1704, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} From afce9f420cbec396301b018b70bd6d72bfbec415 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 02:08:05 +0000 Subject: [PATCH 07/52] Fix pylint import style violations in quantizations_test.py --- tests/unit/quantizations_test.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 662cd23319..765d3de300 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -23,6 +23,7 @@ from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR from flax import nnx from maxtext.layers import moe +from maxtext.layers import linears from maxtext.layers import quantizations from maxtext.kernels.megablox.ops import gmm from maxtext.layers.initializers import nd_dense_init @@ -442,7 +443,6 @@ def test_quantization_fallbacks(self): def test_dense_general_parameter_offload_coverage(self): # Covers parameter_memory_host_offload paths in linears.py - from maxtext.layers import linears dense_layer = linears.DenseGeneral( in_features_shape=8, @@ -479,7 +479,6 @@ def test_configure_quantization_batch_split_schedule(self): def test_moe_gemma4_coverage(self): # Covers GEMMA4 routing and expert scale fusion paths in moe.py - from maxtext.layers import moe config = pyconfig.initialize( [None, get_test_config_path()], From 22b00f6b8d6761a9c8159784d3a2301d32c8dd49 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 02:25:54 +0000 Subject: [PATCH 08/52] Restrict Qwix tiling and group-size overrides to quantized runs in moe.py to prevent VMEM OOM --- src/maxtext/layers/moe.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index c29cb4eeee..28c3dfe5e3 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -1118,7 +1118,7 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments, min(tiling[2], n), ) rhs_inputs = kernel - if self.config.use_qwix_quantization: + if self.config.use_qwix_quantization and self.config.quantization: # Use full contraction for QWIX quantization to allow quantization # fusion (max reduce over contracting dimension). tiling = (tiling[0], k, tiling[2]) @@ -1140,7 +1140,7 @@ def jax_ragged_dot_gmm(inputs, kernel, tiling, group_sizes, expert_assignments, def get_tokamax_group_sizes(group_sizes, inputs, kernel): # TODO (b/491979205) pipeline fsdp ag per repeat fails tokamax gmm - if self.config.use_qwix_quantization or ( + if (self.config.use_qwix_quantization and self.config.quantization) or ( self.config.using_pipeline_parallelism and self.config.pipeline_fsdp_ag_per_repeat ): return group_sizes From 8ac605c63a950d5570f21ae504d4e52c2153b610 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 02:49:32 +0000 Subject: [PATCH 09/52] Fix validation, assertion, and divisibility failures in quantizations_test.py --- tests/unit/quantizations_test.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 765d3de300..4ab2b0a3ca 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -337,7 +337,7 @@ def test_configure_quantization_paths(self): [None, get_test_config_path()], enable_checkpointing=False, quantization="fp8", - use_qwix_quantization=True, + use_qwix_quantization=False, ) quant_fp8 = quantizations.configure_quantization(config_fp8, "train") self.assertIsNotNone(quant_fp8) @@ -346,7 +346,7 @@ def test_configure_quantization_paths(self): [None, get_test_config_path()], enable_checkpointing=False, quantization="nanoo_fp8", - use_qwix_quantization=True, + use_qwix_quantization=False, ) quant_nanoo = quantizations.configure_quantization(config_nanoo, "train") self.assertIsNotNone(quant_nanoo) @@ -364,7 +364,7 @@ def test_configure_quantization_paths(self): [None, get_test_config_path()], enable_checkpointing=False, quantization="te_fp8_delayedscaling", - use_qwix_quantization=True, + use_qwix_quantization=False, ) quant_te = quantizations.configure_quantization(config_te, "train") self.assertIsNotNone(quant_te) @@ -422,9 +422,11 @@ def test_moe_quantization_coverage(self): # Execute a forward pass to cover DenseGeneral.__call__, RoutedMoE.__call__, # sparse_matmul, and the custom quant_einsum wrapper in moe.py - inputs = jnp.ones((2, 4, 8), dtype=jnp.float32) + # Batch size must be a multiple of the mesh axis size (e.g. devices_array.size) to be divisible under FSDP sharding + batch_size = max(4, devices_array.size) + inputs = jnp.ones((batch_size, 4, 8), dtype=jnp.float32) outputs, _, _ = moe_layer(inputs) - self.assertEqual(outputs.shape, (2, 4, 8)) + self.assertEqual(outputs.shape, (batch_size, 4, 8)) def test_quantization_fallbacks(self): # Cover the fallback return None path in _get_quant_config when an unsupported scheme is passed @@ -461,7 +463,7 @@ def test_configure_quantization_batch_split_schedule(self): enable_checkpointing=False, use_batch_split_schedule=True, quantization="fp8_full", - use_qwix_quantization=False, + use_qwix_quantization=True, ) quant = quantizations.configure_quantization(config_bs, "train") self.assertIsInstance(quant, quantizations.QwixQuantization) @@ -472,7 +474,7 @@ def test_configure_quantization_batch_split_schedule(self): use_batch_split_schedule=True, quantization="fp8_full", use_manual_quantization=True, - use_qwix_quantization=False, + use_qwix_quantization=True, ) quant_manual = quantizations.configure_quantization(config_bs_manual, "train") self.assertIsNone(quant_manual) @@ -504,9 +506,11 @@ def test_moe_gemma4_coverage(self): kernel_axes=("expert", "embed_moe", "heads"), rngs=nnx.Rngs(0), ) - inputs = jnp.ones((2, 4, 8), dtype=jnp.float32) + # Batch size must be a multiple of the mesh axis size (e.g. devices_array.size) to be divisible under FSDP sharding + batch_size = max(4, devices_array.size) + inputs = jnp.ones((batch_size, 4, 8), dtype=jnp.float32) outputs, _, _ = moe_layer(inputs) - self.assertEqual(outputs.shape, (2, 4, 8)) + self.assertEqual(outputs.shape, (batch_size, 4, 8)) if __name__ == "__main__": From 9bb5950bd7df758ac44d3590d1d58abba9f56aaa Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 03:33:47 +0000 Subject: [PATCH 10/52] Add @pytest.mark.tpu_only decorator to TPU-specific MoE coverage tests --- tests/unit/quantizations_test.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 4ab2b0a3ca..a8270bc2de 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -386,6 +386,7 @@ def test_configure_kv_quant(self): with self.assertRaises(ValueError): quantizations.configure_kv_quant(config_fail) + @pytest.mark.tpu_only def test_moe_quantization_coverage(self): # Instantiates RoutedMoE on CPU to cover the AQT-free parameter initialization path in moe.py config = pyconfig.initialize( @@ -479,6 +480,7 @@ def test_configure_quantization_batch_split_schedule(self): quant_manual = quantizations.configure_quantization(config_bs_manual, "train") self.assertIsNone(quant_manual) + @pytest.mark.tpu_only def test_moe_gemma4_coverage(self): # Covers GEMMA4 routing and expert scale fusion paths in moe.py From 60a79de5a4ef0c4fb0c017b2875cc3421c178f04 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 04:04:56 +0000 Subject: [PATCH 11/52] Update reference HLO from CI artifact --- tests/utils/reference_hlo_deepseek3.txt | 1502 +++++++-------- tests/utils/reference_hlo_llama3_8b.txt | 2168 ++++++++++----------- tests/utils/reference_hlo_qwen3_1.7b.txt | 2212 +++++++++++----------- 3 files changed, 2941 insertions(+), 2941 deletions(-) diff --git a/tests/utils/reference_hlo_deepseek3.txt b/tests/utils/reference_hlo_deepseek3.txt index de9a31f4e5..ffff9103ad 100644 --- a/tests/utils/reference_hlo_deepseek3.txt +++ b/tests/utils/reference_hlo_deepseek3.txt @@ -220,7 +220,7 @@ StackFrames } %fused_computation.6 (param_0.20: f32[163840,32], param_1.110: s32[1024]) -> f32[512,32] { - %param_0.20 = f32[163840,32]{1,0:T(8,128)S(1)} parameter(0) + %param_0.20 = f32[163840,32]{1,0:T(8,128)} parameter(0) %param_1.110 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.15 = s32[1024]{0:T(1024)} custom-call(%param_1.110), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} %slice.922 = s32[512]{0:T(512)} slice(%custom-call.15), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} @@ -228,7 +228,7 @@ StackFrames %transpose.853 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3326), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} %gather.189 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.20, %transpose.853), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} %transpose.852 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.189), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - ROOT %reshape.3325 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.852), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + ROOT %reshape.3325 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.852), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} } %fused_computation.7 (param_0.23: f32[163840,32], param_1.112: s32[1024]) -> f32[512,32] { @@ -240,7 +240,7 @@ StackFrames %transpose.859 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3334), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} %gather.191 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.23, %transpose.859), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} %transpose.858 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.191), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - ROOT %reshape.3333 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.858), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + ROOT %reshape.3333 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.858), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} } %fused_computation.8 (param_0.26: f32[163840,32], param_1.120: s32[1024]) -> f32[512,32] { @@ -276,7 +276,7 @@ StackFrames %transpose.877 = s32[4096]{0:T(1024)} transpose(%reshape.3358), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} %gather.197 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.32, %transpose.877), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} %transpose.876 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.197), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3357 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.876), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3357 = bf16[4096,512]{1,0:T(8,128)(2,1)} reshape(%transpose.876), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.11 (param_0.35: bf16[4096,512], param_1.128: s32[4096]) -> bf16[4096,512] { @@ -538,14 +538,14 @@ StackFrames ROOT %compare.389 = pred[] compare(%name.16, %name.17), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%called_computation.13 (param_0.4524: s32[256]) -> s32[256] { - %param_0.4524 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2073 = s32[256]{0:T(256)} copy(%param_0.4524), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"1134","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.13 (param_0.4523: s32[256]) -> s32[256] { + %param_0.4523 = s32[256]{0:T(256)} parameter(0) + ROOT %copy.2073 = s32[256]{0:T(256)} copy(%param_0.4523), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"1134","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.13 (param_0.4525: s32[256]) -> s32[256] { - %param_0.4525 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2074.cloned.1 = s32[256]{0:T(256)} call(%param_0.4525), to_apply=%called_computation.13 +%async_computation.13 (param_0.4524: s32[256]) -> s32[256] { + %param_0.4524 = s32[256]{0:T(256)} parameter(0) + ROOT %copy.2074.cloned.1 = s32[256]{0:T(256)} call(%param_0.4524), to_apply=%called_computation.13 }, execution_thread="sparsecore" %region_49.59 (scatter-add.14: s32[], scatter-add.15: s32[]) -> s32[] { @@ -554,33 +554,33 @@ StackFrames ROOT %add.1352 = s32[]{:T(128)S(7)} add(%scatter-add.14, %scatter-add.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.22.clone.clone (param_0.4526: s32[256], param_1.5321: s32[4096], param_2.4492: s32[4096]) -> s32[256] { - %param_0.4526 = s32[256]{0:T(256)} parameter(0) - %param_1.5321 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3923 = s32[4096]{0:T(1024)} reshape(%param_1.5321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} +%fused_computation.22.clone.clone (param_0.4525: s32[256], param_1.5325: s32[4096], param_2.4494: s32[4096]) -> s32[256] { + %param_0.4525 = s32[256]{0:T(256)} parameter(0) + %param_1.5325 = s32[4096]{0:T(1024)} parameter(1) + %reshape.3923 = s32[4096]{0:T(1024)} reshape(%param_1.5325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} %transpose.1100 = s32[4096]{0:T(1024)} transpose(%reshape.3923), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} - %param_2.4492 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3924 = s32[4096]{0:T(1024)} reshape(%param_2.4492), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + %param_2.4494 = s32[4096]{0:T(1024)} parameter(2) + %reshape.3924 = s32[4096]{0:T(1024)} reshape(%param_2.4494), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} %transpose.1101 = s32[4096]{0:T(1024)} transpose(%reshape.3924), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.237 = s32[256]{0:T(256)} scatter(%param_0.4526, %transpose.1100, %transpose.1101), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_49.59, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter-add.237 = s32[256]{0:T(256)} scatter(%param_0.4525, %transpose.1100, %transpose.1101), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_49.59, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.14 (param_0.4527: s32[256], param_1.5322: s32[4096], param_2.4493: s32[4096]) -> s32[256] { - %param_0.4527 = s32[256]{0:T(256)} parameter(0) - %param_1.5322 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4493 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4527, %param_1.5322, %param_2.4493), kind=kCustom, calls=%fused_computation.22.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.14 (param_0.4526: s32[256], param_1.5326: s32[4096], param_2.4495: s32[4096]) -> s32[256] { + %param_0.4526 = s32[256]{0:T(256)} parameter(0) + %param_1.5326 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4495 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4526, %param_1.5326, %param_2.4495), kind=kCustom, calls=%fused_computation.22.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.14 (param_0.4528: s32[256], param_1.5323: s32[4096], param_2.4494: s32[4096]) -> s32[256] { - %param_0.4528 = s32[256]{0:T(256)} parameter(0) - %param_1.5323 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4494 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4528, %param_1.5323, %param_2.4494), to_apply=%called_computation.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} +%async_computation.14 (param_0.4527: s32[256], param_1.5327: s32[4096], param_2.4496: s32[4096]) -> s32[256] { + %param_0.4527 = s32[256]{0:T(256)} parameter(0) + %param_1.5327 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4496 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4527, %param_1.5327, %param_2.4496), to_apply=%called_computation.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation (param_0.84: s32[256], param_1.136: s32[4096], param_2.80: s32[4096], param_3.3083: token[]) -> s32[256] { - %param_3.3083 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation (param_0.84: s32[256], param_1.136: s32[4096], param_2.80: s32[4096], param_3.3085: token[]) -> s32[256] { + %param_3.3085 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.84 = s32[256]{0:T(256)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.136 = s32[4096]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.80 = s32[4096]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -590,22 +590,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.40.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.40.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation (param_0.85: s32[256], param_1.137: s32[4096], param_2.81: s32[4096], param_3.3082: token[]) -> s32[256] { - %param_3.3082 = token[] parameter(3) +%async_computation (param_0.85: s32[256], param_1.137: s32[4096], param_2.81: s32[4096], param_3.3084: token[]) -> s32[256] { + %param_3.3084 = token[] parameter(3) %param_0.85 = s32[256]{0:T(256)} parameter(0) %param_1.137 = s32[4096]{0:T(1024)} parameter(1) %param_2.81 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.2.cloned.1 = s32[256]{0:T(256)} call(%param_0.85, %param_1.137, %param_2.81, %param_3.3082), to_apply=%called_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.2.cloned.1 = s32[256]{0:T(256)} call(%param_0.85, %param_1.137, %param_2.81, %param_3.3084), to_apply=%called_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.15 (param_0.4529: f32[9]) -> f32[9] { - %param_0.4529 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2075 = f32[9]{0:T(128)} copy(%param_0.4529), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.15 (param_0.4528: f32[9]) -> f32[9] { + %param_0.4528 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2075 = f32[9]{0:T(128)} copy(%param_0.4528), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.15 (param_0.4530: f32[9]) -> f32[9] { - %param_0.4530 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2076.cloned.1 = f32[9]{0:T(128)} call(%param_0.4530), to_apply=%called_computation.15 +%async_computation.15 (param_0.4529: f32[9]) -> f32[9] { + %param_0.4529 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2076.cloned.1 = f32[9]{0:T(128)} call(%param_0.4529), to_apply=%called_computation.15 }, execution_thread="sparsecore" %region_61.72 (scatter-add.24: f32[], scatter-add.25: f32[]) -> f32[] { @@ -614,33 +614,33 @@ StackFrames ROOT %add.1358 = f32[]{:T(128)S(7)} add(%scatter-add.24, %scatter-add.25), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.24.clone.clone (param_0.4531: f32[9], param_1.5324: s32[256], param_2.4495: f32[256]) -> f32[9] { - %param_0.4531 = f32[9]{0:T(128)} parameter(0) - %param_1.5324 = s32[256]{0:T(256)} parameter(1) - %reshape.3925 = s32[256]{0:T(256)} reshape(%param_1.5324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} +%fused_computation.24.clone.clone (param_0.4530: f32[9], param_1.5328: s32[256], param_2.4497: f32[256]) -> f32[9] { + %param_0.4530 = f32[9]{0:T(128)} parameter(0) + %param_1.5328 = s32[256]{0:T(256)} parameter(1) + %reshape.3925 = s32[256]{0:T(256)} reshape(%param_1.5328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1102 = s32[256]{0:T(256)} transpose(%reshape.3925), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4495 = f32[256]{0:T(256)} parameter(2) - %reshape.3926 = f32[256]{0:T(256)} reshape(%param_2.4495), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4497 = f32[256]{0:T(256)} parameter(2) + %reshape.3926 = f32[256]{0:T(256)} reshape(%param_2.4497), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1103 = f32[256]{0:T(256)} transpose(%reshape.3926), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.238 = f32[9]{0:T(128)} scatter(%param_0.4531, %transpose.1102, %transpose.1103), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_61.72, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.238 = f32[9]{0:T(128)} scatter(%param_0.4530, %transpose.1102, %transpose.1103), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_61.72, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.16 (param_0.4532: f32[9], param_1.5325: s32[256], param_2.4496: f32[256]) -> f32[9] { - %param_0.4532 = f32[9]{0:T(128)} parameter(0) - %param_1.5325 = s32[256]{0:T(256)} parameter(1) - %param_2.4496 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.41 = f32[9]{0:T(128)} fusion(%param_0.4532, %param_1.5325, %param_2.4496), kind=kCustom, calls=%fused_computation.24.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.16 (param_0.4531: f32[9], param_1.5329: s32[256], param_2.4498: f32[256]) -> f32[9] { + %param_0.4531 = f32[9]{0:T(128)} parameter(0) + %param_1.5329 = s32[256]{0:T(256)} parameter(1) + %param_2.4498 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.41 = f32[9]{0:T(128)} fusion(%param_0.4531, %param_1.5329, %param_2.4498), kind=kCustom, calls=%fused_computation.24.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.16 (param_0.4533: f32[9], param_1.5326: s32[256], param_2.4497: f32[256]) -> f32[9] { - %param_0.4533 = f32[9]{0:T(128)} parameter(0) - %param_1.5326 = s32[256]{0:T(256)} parameter(1) - %param_2.4497 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.42.cloned.1 = f32[9]{0:T(128)} call(%param_0.4533, %param_1.5326, %param_2.4497), to_apply=%called_computation.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +%async_computation.16 (param_0.4532: f32[9], param_1.5330: s32[256], param_2.4499: f32[256]) -> f32[9] { + %param_0.4532 = f32[9]{0:T(128)} parameter(0) + %param_1.5330 = s32[256]{0:T(256)} parameter(1) + %param_2.4499 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.42.cloned.1 = f32[9]{0:T(128)} call(%param_0.4532, %param_1.5330, %param_2.4499), to_apply=%called_computation.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.1 (param_0.87: f32[9], param_1.139: s32[256], param_2.83: f32[256], param_3.3097: token[]) -> f32[9] { - %param_3.3097 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.1 (param_0.87: f32[9], param_1.139: s32[256], param_2.83: f32[256], param_3.3099: token[]) -> f32[9] { + %param_3.3099 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.87 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.139 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.83 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -650,22 +650,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.42.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.42.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.1 (param_0.88: f32[9], param_1.140: s32[256], param_2.84: f32[256], param_3.3096: token[]) -> f32[9] { - %param_3.3096 = token[] parameter(3) +%async_computation.1 (param_0.88: f32[9], param_1.140: s32[256], param_2.84: f32[256], param_3.3098: token[]) -> f32[9] { + %param_3.3098 = token[] parameter(3) %param_0.88 = f32[9]{0:T(128)} parameter(0) %param_1.140 = s32[256]{0:T(256)} parameter(1) %param_2.84 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.5.cloned.1 = f32[9]{0:T(128)} call(%param_0.88, %param_1.140, %param_2.84, %param_3.3096), to_apply=%called_computation.1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.5.cloned.1 = f32[9]{0:T(128)} call(%param_0.88, %param_1.140, %param_2.84, %param_3.3098), to_apply=%called_computation.1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.17 (param_0.4534: s32[263]) -> s32[263] { - %param_0.4534 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2077 = s32[263]{0:T(512)} copy(%param_0.4534), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.17 (param_0.4533: s32[263]) -> s32[263] { + %param_0.4533 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2077 = s32[263]{0:T(512)} copy(%param_0.4533), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.17 (param_0.4535: s32[263]) -> s32[263] { - %param_0.4535 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2078.cloned.1 = s32[263]{0:T(512)} call(%param_0.4535), to_apply=%called_computation.17 +%async_computation.17 (param_0.4534: s32[263]) -> s32[263] { + %param_0.4534 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2078.cloned.1 = s32[263]{0:T(512)} call(%param_0.4534), to_apply=%called_computation.17 }, execution_thread="sparsecore" %region_63.74 (scatter-add.28: s32[], scatter-add.29: s32[]) -> s32[] { @@ -674,33 +674,33 @@ StackFrames ROOT %add.1359 = s32[]{:T(128)S(7)} add(%scatter-add.28, %scatter-add.29), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.25.clone.clone (param_0.4536: s32[263], param_1.5327: s32[8], param_2.4498: s32[8]) -> s32[263] { - %param_0.4536 = s32[263]{0:T(512)} parameter(0) - %param_1.5327 = s32[8]{0:T(128)} parameter(1) - %reshape.3927 = s32[8]{0:T(128)} reshape(%param_1.5327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} +%fused_computation.25.clone.clone (param_0.4535: s32[263], param_1.5331: s32[8], param_2.4500: s32[8]) -> s32[263] { + %param_0.4535 = s32[263]{0:T(512)} parameter(0) + %param_1.5331 = s32[8]{0:T(128)} parameter(1) + %reshape.3927 = s32[8]{0:T(128)} reshape(%param_1.5331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} %transpose.1104 = s32[8]{0:T(128)} transpose(%reshape.3927), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4498 = s32[8]{0:T(128)} parameter(2) - %reshape.3928 = s32[8]{0:T(128)} reshape(%param_2.4498), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %param_2.4500 = s32[8]{0:T(128)} parameter(2) + %reshape.3928 = s32[8]{0:T(128)} reshape(%param_2.4500), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} %transpose.1105 = s32[8]{0:T(128)} transpose(%reshape.3928), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.239 = s32[263]{0:T(512)} scatter(%param_0.4536, %transpose.1104, %transpose.1105), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_63.74, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.239 = s32[263]{0:T(512)} scatter(%param_0.4535, %transpose.1104, %transpose.1105), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_63.74, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.18 (param_0.4537: s32[263], param_1.5328: s32[8], param_2.4499: s32[8]) -> s32[263] { - %param_0.4537 = s32[263]{0:T(512)} parameter(0) - %param_1.5328 = s32[8]{0:T(128)} parameter(1) - %param_2.4499 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.43 = s32[263]{0:T(512)} fusion(%param_0.4537, %param_1.5328, %param_2.4499), kind=kCustom, calls=%fused_computation.25.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.18 (param_0.4536: s32[263], param_1.5332: s32[8], param_2.4501: s32[8]) -> s32[263] { + %param_0.4536 = s32[263]{0:T(512)} parameter(0) + %param_1.5332 = s32[8]{0:T(128)} parameter(1) + %param_2.4501 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.43 = s32[263]{0:T(512)} fusion(%param_0.4536, %param_1.5332, %param_2.4501), kind=kCustom, calls=%fused_computation.25.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.18 (param_0.4538: s32[263], param_1.5329: s32[8], param_2.4500: s32[8]) -> s32[263] { - %param_0.4538 = s32[263]{0:T(512)} parameter(0) - %param_1.5329 = s32[8]{0:T(128)} parameter(1) - %param_2.4500 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.44.cloned.1 = s32[263]{0:T(512)} call(%param_0.4538, %param_1.5329, %param_2.4500), to_apply=%called_computation.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +%async_computation.18 (param_0.4537: s32[263], param_1.5333: s32[8], param_2.4502: s32[8]) -> s32[263] { + %param_0.4537 = s32[263]{0:T(512)} parameter(0) + %param_1.5333 = s32[8]{0:T(128)} parameter(1) + %param_2.4502 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.44.cloned.1 = s32[263]{0:T(512)} call(%param_0.4537, %param_1.5333, %param_2.4502), to_apply=%called_computation.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.2 (param_0.90: s32[263], param_1.142: s32[8], param_2.86: s32[8], param_3.3103: token[]) -> s32[263] { - %param_3.3103 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.2 (param_0.90: s32[263], param_1.142: s32[8], param_2.86: s32[8], param_3.3105: token[]) -> s32[263] { + %param_3.3105 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.90 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.142 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.86 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -710,22 +710,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.44.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.44.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.2 (param_0.91: s32[263], param_1.143: s32[8], param_2.87: s32[8], param_3.3102: token[]) -> s32[263] { - %param_3.3102 = token[] parameter(3) +%async_computation.2 (param_0.91: s32[263], param_1.143: s32[8], param_2.87: s32[8], param_3.3104: token[]) -> s32[263] { + %param_3.3104 = token[] parameter(3) %param_0.91 = s32[263]{0:T(512)} parameter(0) %param_1.143 = s32[8]{0:T(128)} parameter(1) %param_2.87 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.8.cloned.1 = s32[263]{0:T(512)} call(%param_0.91, %param_1.143, %param_2.87, %param_3.3102), to_apply=%called_computation.2, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.8.cloned.1 = s32[263]{0:T(512)} call(%param_0.91, %param_1.143, %param_2.87, %param_3.3104), to_apply=%called_computation.2, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.19 (param_0.4539: s32[263]) -> s32[263] { - %param_0.4539 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2079 = s32[263]{0:T(512)} copy(%param_0.4539), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.19 (param_0.4538: s32[263]) -> s32[263] { + %param_0.4538 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2079 = s32[263]{0:T(512)} copy(%param_0.4538), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.19 (param_0.4540: s32[263]) -> s32[263] { - %param_0.4540 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2080.cloned.1 = s32[263]{0:T(512)} call(%param_0.4540), to_apply=%called_computation.19 +%async_computation.19 (param_0.4539: s32[263]) -> s32[263] { + %param_0.4539 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2080.cloned.1 = s32[263]{0:T(512)} call(%param_0.4539), to_apply=%called_computation.19 }, execution_thread="sparsecore" %region_73.86.clone (scatter-add.163: s32[], scatter-add.164: s32[]) -> s32[] { @@ -734,33 +734,33 @@ StackFrames ROOT %add.2474 = s32[]{:T(128)S(7)} add(%scatter-add.163, %scatter-add.164), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.26.clone.clone (param_0.4541: s32[263], param_1.5330: s32[256], param_2.4501: s32[256]) -> s32[263] { - %param_0.4541 = s32[263]{0:T(512)} parameter(0) - %param_1.5330 = s32[256]{0:T(256)} parameter(1) - %reshape.3929 = s32[256]{0:T(256)} reshape(%param_1.5330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} +%fused_computation.26.clone.clone (param_0.4540: s32[263], param_1.5334: s32[256], param_2.4503: s32[256]) -> s32[263] { + %param_0.4540 = s32[263]{0:T(512)} parameter(0) + %param_1.5334 = s32[256]{0:T(256)} parameter(1) + %reshape.3929 = s32[256]{0:T(256)} reshape(%param_1.5334), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} %transpose.1106 = s32[256]{0:T(256)} transpose(%reshape.3929), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4501 = s32[256]{0:T(256)} parameter(2) - %reshape.3930 = s32[256]{0:T(256)} reshape(%param_2.4501), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4503 = s32[256]{0:T(256)} parameter(2) + %reshape.3930 = s32[256]{0:T(256)} reshape(%param_2.4503), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1107 = s32[256]{0:T(256)} transpose(%reshape.3930), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.240 = s32[263]{0:T(512)} scatter(%param_0.4541, %transpose.1106, %transpose.1107), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_73.86.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.240 = s32[263]{0:T(512)} scatter(%param_0.4540, %transpose.1106, %transpose.1107), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_73.86.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.20 (param_0.4542: s32[263], param_1.5331: s32[256], param_2.4502: s32[256]) -> s32[263] { - %param_0.4542 = s32[263]{0:T(512)} parameter(0) - %param_1.5331 = s32[256]{0:T(256)} parameter(1) - %param_2.4502 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.45 = s32[263]{0:T(512)} fusion(%param_0.4542, %param_1.5331, %param_2.4502), kind=kCustom, calls=%fused_computation.26.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.20 (param_0.4541: s32[263], param_1.5335: s32[256], param_2.4504: s32[256]) -> s32[263] { + %param_0.4541 = s32[263]{0:T(512)} parameter(0) + %param_1.5335 = s32[256]{0:T(256)} parameter(1) + %param_2.4504 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.45 = s32[263]{0:T(512)} fusion(%param_0.4541, %param_1.5335, %param_2.4504), kind=kCustom, calls=%fused_computation.26.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.20 (param_0.4543: s32[263], param_1.5332: s32[256], param_2.4503: s32[256]) -> s32[263] { - %param_0.4543 = s32[263]{0:T(512)} parameter(0) - %param_1.5332 = s32[256]{0:T(256)} parameter(1) - %param_2.4503 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.46.cloned.1 = s32[263]{0:T(512)} call(%param_0.4543, %param_1.5332, %param_2.4503), to_apply=%called_computation.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +%async_computation.20 (param_0.4542: s32[263], param_1.5336: s32[256], param_2.4505: s32[256]) -> s32[263] { + %param_0.4542 = s32[263]{0:T(512)} parameter(0) + %param_1.5336 = s32[256]{0:T(256)} parameter(1) + %param_2.4505 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.46.cloned.1 = s32[263]{0:T(512)} call(%param_0.4542, %param_1.5336, %param_2.4505), to_apply=%called_computation.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.3 (param_0.93: s32[263], param_1.145: s32[256], param_2.89: s32[256], param_3.3089: token[]) -> s32[263] { - %param_3.3089 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.3 (param_0.93: s32[263], param_1.145: s32[256], param_2.89: s32[256], param_3.3091: token[]) -> s32[263] { + %param_3.3091 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.93 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.145 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.89 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -770,22 +770,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.46.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.46.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.3 (param_0.94: s32[263], param_1.146: s32[256], param_2.90: s32[256], param_3.3088: token[]) -> s32[263] { - %param_3.3088 = token[] parameter(3) +%async_computation.3 (param_0.94: s32[263], param_1.146: s32[256], param_2.90: s32[256], param_3.3090: token[]) -> s32[263] { + %param_3.3090 = token[] parameter(3) %param_0.94 = s32[263]{0:T(512)} parameter(0) %param_1.146 = s32[256]{0:T(256)} parameter(1) %param_2.90 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.11.cloned.1 = s32[263]{0:T(512)} call(%param_0.94, %param_1.146, %param_2.90, %param_3.3088), to_apply=%called_computation.3, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.11.cloned.1 = s32[263]{0:T(512)} call(%param_0.94, %param_1.146, %param_2.90, %param_3.3090), to_apply=%called_computation.3, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.21 (param_0.4544: f32[9]) -> f32[9] { - %param_0.4544 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2081 = f32[9]{0:T(128)} copy(%param_0.4544), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.21 (param_0.4543: f32[9]) -> f32[9] { + %param_0.4543 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2081 = f32[9]{0:T(128)} copy(%param_0.4543), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.21 (param_0.4545: f32[9]) -> f32[9] { - %param_0.4545 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2082.cloned.1 = f32[9]{0:T(128)} call(%param_0.4545), to_apply=%called_computation.21 +%async_computation.21 (param_0.4544: f32[9]) -> f32[9] { + %param_0.4544 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2082.cloned.1 = f32[9]{0:T(128)} call(%param_0.4544), to_apply=%called_computation.21 }, execution_thread="sparsecore" %region_79.95.clone (scatter-add.167: f32[], scatter-add.168: f32[]) -> f32[] { @@ -794,33 +794,33 @@ StackFrames ROOT %add.2476 = f32[]{:T(128)S(7)} add(%scatter-add.167, %scatter-add.168), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.27.clone.clone (param_0.4546: f32[9], param_1.5333: s32[256], param_2.4504: f32[256]) -> f32[9] { - %param_0.4546 = f32[9]{0:T(128)} parameter(0) - %param_1.5333 = s32[256]{0:T(256)} parameter(1) - %reshape.3931 = s32[256]{0:T(256)} reshape(%param_1.5333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} +%fused_computation.27.clone.clone (param_0.4545: f32[9], param_1.5337: s32[256], param_2.4506: f32[256]) -> f32[9] { + %param_0.4545 = f32[9]{0:T(128)} parameter(0) + %param_1.5337 = s32[256]{0:T(256)} parameter(1) + %reshape.3931 = s32[256]{0:T(256)} reshape(%param_1.5337), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1108 = s32[256]{0:T(256)} transpose(%reshape.3931), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4504 = f32[256]{0:T(256)} parameter(2) - %reshape.3932 = f32[256]{0:T(256)} reshape(%param_2.4504), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4506 = f32[256]{0:T(256)} parameter(2) + %reshape.3932 = f32[256]{0:T(256)} reshape(%param_2.4506), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1109 = f32[256]{0:T(256)} transpose(%reshape.3932), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.241 = f32[9]{0:T(128)} scatter(%param_0.4546, %transpose.1108, %transpose.1109), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_79.95.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.241 = f32[9]{0:T(128)} scatter(%param_0.4545, %transpose.1108, %transpose.1109), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_79.95.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.22 (param_0.4547: f32[9], param_1.5334: s32[256], param_2.4505: f32[256]) -> f32[9] { - %param_0.4547 = f32[9]{0:T(128)} parameter(0) - %param_1.5334 = s32[256]{0:T(256)} parameter(1) - %param_2.4505 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.47 = f32[9]{0:T(128)} fusion(%param_0.4547, %param_1.5334, %param_2.4505), kind=kCustom, calls=%fused_computation.27.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.22 (param_0.4546: f32[9], param_1.5338: s32[256], param_2.4507: f32[256]) -> f32[9] { + %param_0.4546 = f32[9]{0:T(128)} parameter(0) + %param_1.5338 = s32[256]{0:T(256)} parameter(1) + %param_2.4507 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.47 = f32[9]{0:T(128)} fusion(%param_0.4546, %param_1.5338, %param_2.4507), kind=kCustom, calls=%fused_computation.27.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.22 (param_0.4548: f32[9], param_1.5335: s32[256], param_2.4506: f32[256]) -> f32[9] { - %param_0.4548 = f32[9]{0:T(128)} parameter(0) - %param_1.5335 = s32[256]{0:T(256)} parameter(1) - %param_2.4506 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.48.cloned.1 = f32[9]{0:T(128)} call(%param_0.4548, %param_1.5335, %param_2.4506), to_apply=%called_computation.22, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +%async_computation.22 (param_0.4547: f32[9], param_1.5339: s32[256], param_2.4508: f32[256]) -> f32[9] { + %param_0.4547 = f32[9]{0:T(128)} parameter(0) + %param_1.5339 = s32[256]{0:T(256)} parameter(1) + %param_2.4508 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.48.cloned.1 = f32[9]{0:T(128)} call(%param_0.4547, %param_1.5339, %param_2.4508), to_apply=%called_computation.22, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.4 (param_0.96: f32[9], param_1.148: s32[256], param_2.92: f32[256], param_3.3095: token[]) -> f32[9] { - %param_3.3095 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.4 (param_0.96: f32[9], param_1.148: s32[256], param_2.92: f32[256], param_3.3097: token[]) -> f32[9] { + %param_3.3097 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.96 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.148 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.92 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -830,22 +830,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.48.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.48.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.4 (param_0.97: f32[9], param_1.149: s32[256], param_2.93: f32[256], param_3.3094: token[]) -> f32[9] { - %param_3.3094 = token[] parameter(3) +%async_computation.4 (param_0.97: f32[9], param_1.149: s32[256], param_2.93: f32[256], param_3.3096: token[]) -> f32[9] { + %param_3.3096 = token[] parameter(3) %param_0.97 = f32[9]{0:T(128)} parameter(0) %param_1.149 = s32[256]{0:T(256)} parameter(1) %param_2.93 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.14.cloned.1 = f32[9]{0:T(128)} call(%param_0.97, %param_1.149, %param_2.93, %param_3.3094), to_apply=%called_computation.4, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.14.cloned.1 = f32[9]{0:T(128)} call(%param_0.97, %param_1.149, %param_2.93, %param_3.3096), to_apply=%called_computation.4, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.23 (param_0.4549: s32[263]) -> s32[263] { - %param_0.4549 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2083 = s32[263]{0:T(512)} copy(%param_0.4549), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.23 (param_0.4548: s32[263]) -> s32[263] { + %param_0.4548 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2083 = s32[263]{0:T(512)} copy(%param_0.4548), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.23 (param_0.4550: s32[263]) -> s32[263] { - %param_0.4550 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2084.cloned.1 = s32[263]{0:T(512)} call(%param_0.4550), to_apply=%called_computation.23 +%async_computation.23 (param_0.4549: s32[263]) -> s32[263] { + %param_0.4549 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2084.cloned.1 = s32[263]{0:T(512)} call(%param_0.4549), to_apply=%called_computation.23 }, execution_thread="sparsecore" %region_81.97.clone (scatter-add.171: s32[], scatter-add.172: s32[]) -> s32[] { @@ -854,33 +854,33 @@ StackFrames ROOT %add.2478 = s32[]{:T(128)S(7)} add(%scatter-add.171, %scatter-add.172), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.28.clone.clone (param_0.4551: s32[263], param_1.5336: s32[8], param_2.4507: s32[8]) -> s32[263] { - %param_0.4551 = s32[263]{0:T(512)} parameter(0) - %param_1.5336 = s32[8]{0:T(128)} parameter(1) - %reshape.3933 = s32[8]{0:T(128)} reshape(%param_1.5336), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} +%fused_computation.28.clone.clone (param_0.4550: s32[263], param_1.5340: s32[8], param_2.4509: s32[8]) -> s32[263] { + %param_0.4550 = s32[263]{0:T(512)} parameter(0) + %param_1.5340 = s32[8]{0:T(128)} parameter(1) + %reshape.3933 = s32[8]{0:T(128)} reshape(%param_1.5340), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} %transpose.1110 = s32[8]{0:T(128)} transpose(%reshape.3933), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4507 = s32[8]{0:T(128)} parameter(2) - %reshape.3934 = s32[8]{0:T(128)} reshape(%param_2.4507), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %param_2.4509 = s32[8]{0:T(128)} parameter(2) + %reshape.3934 = s32[8]{0:T(128)} reshape(%param_2.4509), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} %transpose.1111 = s32[8]{0:T(128)} transpose(%reshape.3934), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.242 = s32[263]{0:T(512)} scatter(%param_0.4551, %transpose.1110, %transpose.1111), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_81.97.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.242 = s32[263]{0:T(512)} scatter(%param_0.4550, %transpose.1110, %transpose.1111), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_81.97.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.24 (param_0.4552: s32[263], param_1.5337: s32[8], param_2.4508: s32[8]) -> s32[263] { - %param_0.4552 = s32[263]{0:T(512)} parameter(0) - %param_1.5337 = s32[8]{0:T(128)} parameter(1) - %param_2.4508 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.49 = s32[263]{0:T(512)} fusion(%param_0.4552, %param_1.5337, %param_2.4508), kind=kCustom, calls=%fused_computation.28.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.24 (param_0.4551: s32[263], param_1.5341: s32[8], param_2.4510: s32[8]) -> s32[263] { + %param_0.4551 = s32[263]{0:T(512)} parameter(0) + %param_1.5341 = s32[8]{0:T(128)} parameter(1) + %param_2.4510 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.49 = s32[263]{0:T(512)} fusion(%param_0.4551, %param_1.5341, %param_2.4510), kind=kCustom, calls=%fused_computation.28.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.24 (param_0.4553: s32[263], param_1.5338: s32[8], param_2.4509: s32[8]) -> s32[263] { - %param_0.4553 = s32[263]{0:T(512)} parameter(0) - %param_1.5338 = s32[8]{0:T(128)} parameter(1) - %param_2.4509 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.50.cloned.1 = s32[263]{0:T(512)} call(%param_0.4553, %param_1.5338, %param_2.4509), to_apply=%called_computation.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +%async_computation.24 (param_0.4552: s32[263], param_1.5342: s32[8], param_2.4511: s32[8]) -> s32[263] { + %param_0.4552 = s32[263]{0:T(512)} parameter(0) + %param_1.5342 = s32[8]{0:T(128)} parameter(1) + %param_2.4511 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.50.cloned.1 = s32[263]{0:T(512)} call(%param_0.4552, %param_1.5342, %param_2.4511), to_apply=%called_computation.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.5 (param_0.99: s32[263], param_1.151: s32[8], param_2.95: s32[8], param_3.3105: token[]) -> s32[263] { - %param_3.3105 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.5 (param_0.99: s32[263], param_1.151: s32[8], param_2.95: s32[8], param_3.3107: token[]) -> s32[263] { + %param_3.3107 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.99 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.151 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.95 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -890,22 +890,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.50.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.50.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.5 (param_0.100: s32[263], param_1.152: s32[8], param_2.96: s32[8], param_3.3104: token[]) -> s32[263] { - %param_3.3104 = token[] parameter(3) +%async_computation.5 (param_0.100: s32[263], param_1.152: s32[8], param_2.96: s32[8], param_3.3106: token[]) -> s32[263] { + %param_3.3106 = token[] parameter(3) %param_0.100 = s32[263]{0:T(512)} parameter(0) %param_1.152 = s32[8]{0:T(128)} parameter(1) %param_2.96 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.17.cloned.1 = s32[263]{0:T(512)} call(%param_0.100, %param_1.152, %param_2.96, %param_3.3104), to_apply=%called_computation.5, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.17.cloned.1 = s32[263]{0:T(512)} call(%param_0.100, %param_1.152, %param_2.96, %param_3.3106), to_apply=%called_computation.5, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.25 (param_0.4554: s32[263]) -> s32[263] { - %param_0.4554 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2085 = s32[263]{0:T(512)} copy(%param_0.4554), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.25 (param_0.4553: s32[263]) -> s32[263] { + %param_0.4553 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2085 = s32[263]{0:T(512)} copy(%param_0.4553), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.25 (param_0.4555: s32[263]) -> s32[263] { - %param_0.4555 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2086.cloned.1 = s32[263]{0:T(512)} call(%param_0.4555), to_apply=%called_computation.25 +%async_computation.25 (param_0.4554: s32[263]) -> s32[263] { + %param_0.4554 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2086.cloned.1 = s32[263]{0:T(512)} call(%param_0.4554), to_apply=%called_computation.25 }, execution_thread="sparsecore" %region_96.114 (scatter-add.48: s32[], scatter-add.49: s32[]) -> s32[] { @@ -914,33 +914,33 @@ StackFrames ROOT %add.1396 = s32[]{:T(128)S(7)} add(%scatter-add.48, %scatter-add.49), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.29.clone.clone (param_0.4556: s32[263], param_1.5339: s32[256], param_2.4510: s32[256]) -> s32[263] { - %param_0.4556 = s32[263]{0:T(512)} parameter(0) - %param_1.5339 = s32[256]{0:T(256)} parameter(1) - %reshape.3935 = s32[256]{0:T(256)} reshape(%param_1.5339), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} +%fused_computation.29.clone.clone (param_0.4555: s32[263], param_1.5343: s32[256], param_2.4512: s32[256]) -> s32[263] { + %param_0.4555 = s32[263]{0:T(512)} parameter(0) + %param_1.5343 = s32[256]{0:T(256)} parameter(1) + %reshape.3935 = s32[256]{0:T(256)} reshape(%param_1.5343), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} %transpose.1112 = s32[256]{0:T(256)} transpose(%reshape.3935), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %param_2.4510 = s32[256]{0:T(256)} parameter(2) - %reshape.3936 = s32[256]{0:T(256)} reshape(%param_2.4510), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4512 = s32[256]{0:T(256)} parameter(2) + %reshape.3936 = s32[256]{0:T(256)} reshape(%param_2.4512), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1113 = s32[256]{0:T(256)} transpose(%reshape.3936), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.243 = s32[263]{0:T(512)} scatter(%param_0.4556, %transpose.1112, %transpose.1113), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_96.114, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.243 = s32[263]{0:T(512)} scatter(%param_0.4555, %transpose.1112, %transpose.1113), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_96.114, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.26 (param_0.4557: s32[263], param_1.5340: s32[256], param_2.4511: s32[256]) -> s32[263] { - %param_0.4557 = s32[263]{0:T(512)} parameter(0) - %param_1.5340 = s32[256]{0:T(256)} parameter(1) - %param_2.4511 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.51 = s32[263]{0:T(512)} fusion(%param_0.4557, %param_1.5340, %param_2.4511), kind=kCustom, calls=%fused_computation.29.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.26 (param_0.4556: s32[263], param_1.5344: s32[256], param_2.4513: s32[256]) -> s32[263] { + %param_0.4556 = s32[263]{0:T(512)} parameter(0) + %param_1.5344 = s32[256]{0:T(256)} parameter(1) + %param_2.4513 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.51 = s32[263]{0:T(512)} fusion(%param_0.4556, %param_1.5344, %param_2.4513), kind=kCustom, calls=%fused_computation.29.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.26 (param_0.4558: s32[263], param_1.5341: s32[256], param_2.4512: s32[256]) -> s32[263] { - %param_0.4558 = s32[263]{0:T(512)} parameter(0) - %param_1.5341 = s32[256]{0:T(256)} parameter(1) - %param_2.4512 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.52.cloned.1 = s32[263]{0:T(512)} call(%param_0.4558, %param_1.5341, %param_2.4512), to_apply=%called_computation.26, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} +%async_computation.26 (param_0.4557: s32[263], param_1.5345: s32[256], param_2.4514: s32[256]) -> s32[263] { + %param_0.4557 = s32[263]{0:T(512)} parameter(0) + %param_1.5345 = s32[256]{0:T(256)} parameter(1) + %param_2.4514 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.52.cloned.1 = s32[263]{0:T(512)} call(%param_0.4557, %param_1.5345, %param_2.4514), to_apply=%called_computation.26, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.6 (param_0.102: s32[263], param_1.154: s32[256], param_2.98: s32[256], param_3.3091: token[]) -> s32[263] { - %param_3.3091 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.6 (param_0.102: s32[263], param_1.154: s32[256], param_2.98: s32[256], param_3.3093: token[]) -> s32[263] { + %param_3.3093 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.102 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.154 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.98 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -950,12 +950,12 @@ StackFrames ROOT %scatter_offload_custom_fusion.52.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.52.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.6 (param_0.103: s32[263], param_1.155: s32[256], param_2.99: s32[256], param_3.3090: token[]) -> s32[263] { - %param_3.3090 = token[] parameter(3) +%async_computation.6 (param_0.103: s32[263], param_1.155: s32[256], param_2.99: s32[256], param_3.3092: token[]) -> s32[263] { + %param_3.3092 = token[] parameter(3) %param_0.103 = s32[263]{0:T(512)} parameter(0) %param_1.155 = s32[256]{0:T(256)} parameter(1) %param_2.99 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.20.cloned.1 = s32[263]{0:T(512)} call(%param_0.103, %param_1.155, %param_2.99, %param_3.3090), to_apply=%called_computation.6, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.20.cloned.1 = s32[263]{0:T(512)} call(%param_0.103, %param_1.155, %param_2.99, %param_3.3092), to_apply=%called_computation.6, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_102.120 (scatter-add.52: f32[], scatter-add.53: f32[]) -> f32[] { @@ -964,33 +964,33 @@ StackFrames ROOT %add.1399 = f32[]{:T(128)S(7)} add(%scatter-add.52, %scatter-add.53), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.30.clone.clone (param_0.4561: f32[9], param_1.5342: s32[256], param_2.4513: f32[256]) -> f32[9] { - %param_0.4561 = f32[9]{0:T(128)} parameter(0) - %param_1.5342 = s32[256]{0:T(256)} parameter(1) - %reshape.3937 = s32[256]{0:T(256)} reshape(%param_1.5342), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} +%fused_computation.30.clone.clone (param_0.4560: f32[9], param_1.5346: s32[256], param_2.4515: f32[256]) -> f32[9] { + %param_0.4560 = f32[9]{0:T(128)} parameter(0) + %param_1.5346 = s32[256]{0:T(256)} parameter(1) + %reshape.3937 = s32[256]{0:T(256)} reshape(%param_1.5346), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1114 = s32[256]{0:T(256)} transpose(%reshape.3937), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4513 = f32[256]{0:T(256)} parameter(2) - %reshape.3938 = f32[256]{0:T(256)} reshape(%param_2.4513), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4515 = f32[256]{0:T(256)} parameter(2) + %reshape.3938 = f32[256]{0:T(256)} reshape(%param_2.4515), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1115 = f32[256]{0:T(256)} transpose(%reshape.3938), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.244 = f32[9]{0:T(128)} scatter(%param_0.4561, %transpose.1114, %transpose.1115), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_102.120, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.244 = f32[9]{0:T(128)} scatter(%param_0.4560, %transpose.1114, %transpose.1115), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_102.120, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.28 (param_0.4562: f32[9], param_1.5343: s32[256], param_2.4514: f32[256]) -> f32[9] { - %param_0.4562 = f32[9]{0:T(128)} parameter(0) - %param_1.5343 = s32[256]{0:T(256)} parameter(1) - %param_2.4514 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.53 = f32[9]{0:T(128)} fusion(%param_0.4562, %param_1.5343, %param_2.4514), kind=kCustom, calls=%fused_computation.30.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.28 (param_0.4561: f32[9], param_1.5347: s32[256], param_2.4516: f32[256]) -> f32[9] { + %param_0.4561 = f32[9]{0:T(128)} parameter(0) + %param_1.5347 = s32[256]{0:T(256)} parameter(1) + %param_2.4516 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.53 = f32[9]{0:T(128)} fusion(%param_0.4561, %param_1.5347, %param_2.4516), kind=kCustom, calls=%fused_computation.30.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.28 (param_0.4563: f32[9], param_1.5344: s32[256], param_2.4515: f32[256]) -> f32[9] { - %param_0.4563 = f32[9]{0:T(128)} parameter(0) - %param_1.5344 = s32[256]{0:T(256)} parameter(1) - %param_2.4515 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.54.cloned.1 = f32[9]{0:T(128)} call(%param_0.4563, %param_1.5344, %param_2.4515), to_apply=%called_computation.28, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} +%async_computation.28 (param_0.4562: f32[9], param_1.5348: s32[256], param_2.4517: f32[256]) -> f32[9] { + %param_0.4562 = f32[9]{0:T(128)} parameter(0) + %param_1.5348 = s32[256]{0:T(256)} parameter(1) + %param_2.4517 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.54.cloned.1 = f32[9]{0:T(128)} call(%param_0.4562, %param_1.5348, %param_2.4517), to_apply=%called_computation.28, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.7 (param_0.105: f32[9], param_1.157: s32[256], param_2.101: f32[256], param_3.3099: token[]) -> f32[9] { - %param_3.3099 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.7 (param_0.105: f32[9], param_1.157: s32[256], param_2.101: f32[256], param_3.3101: token[]) -> f32[9] { + %param_3.3101 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.105 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.157 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.101 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -998,12 +998,12 @@ StackFrames ROOT %scatter_offload_custom_fusion.54.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.54.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.7 (param_0.106: f32[9], param_1.158: s32[256], param_2.102: f32[256], param_3.3098: token[]) -> f32[9] { - %param_3.3098 = token[] parameter(3) +%async_computation.7 (param_0.106: f32[9], param_1.158: s32[256], param_2.102: f32[256], param_3.3100: token[]) -> f32[9] { + %param_3.3100 = token[] parameter(3) %param_0.106 = f32[9]{0:T(128)} parameter(0) %param_1.158 = s32[256]{0:T(256)} parameter(1) %param_2.102 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.23.cloned.1 = f32[9]{0:T(128)} call(%param_0.106, %param_1.158, %param_2.102, %param_3.3098), to_apply=%called_computation.7, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.23.cloned.1 = f32[9]{0:T(128)} call(%param_0.106, %param_1.158, %param_2.102, %param_3.3100), to_apply=%called_computation.7, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_104.122 (scatter-add.83: s32[], scatter-add.84: s32[]) -> s32[] { @@ -1012,33 +1012,33 @@ StackFrames ROOT %add.1400 = s32[]{:T(128)S(7)} add(%scatter-add.83, %scatter-add.84), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.31.clone.clone (param_0.4566: s32[263], param_1.5345: s32[8], param_2.4516: s32[8]) -> s32[263] { - %param_0.4566 = s32[263]{0:T(512)} parameter(0) - %param_1.5345 = s32[8]{0:T(128)} parameter(1) - %reshape.3939 = s32[8]{0:T(128)} reshape(%param_1.5345), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} +%fused_computation.31.clone.clone (param_0.4565: s32[263], param_1.5349: s32[8], param_2.4518: s32[8]) -> s32[263] { + %param_0.4565 = s32[263]{0:T(512)} parameter(0) + %param_1.5349 = s32[8]{0:T(128)} parameter(1) + %reshape.3939 = s32[8]{0:T(128)} reshape(%param_1.5349), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} %transpose.1116 = s32[8]{0:T(128)} transpose(%reshape.3939), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %param_2.4516 = s32[8]{0:T(128)} parameter(2) - %reshape.3940 = s32[8]{0:T(128)} reshape(%param_2.4516), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %param_2.4518 = s32[8]{0:T(128)} parameter(2) + %reshape.3940 = s32[8]{0:T(128)} reshape(%param_2.4518), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} %transpose.1117 = s32[8]{0:T(128)} transpose(%reshape.3940), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.245 = s32[263]{0:T(512)} scatter(%param_0.4566, %transpose.1116, %transpose.1117), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_104.122, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.245 = s32[263]{0:T(512)} scatter(%param_0.4565, %transpose.1116, %transpose.1117), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_104.122, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.30 (param_0.4567: s32[263], param_1.5346: s32[8], param_2.4517: s32[8]) -> s32[263] { - %param_0.4567 = s32[263]{0:T(512)} parameter(0) - %param_1.5346 = s32[8]{0:T(128)} parameter(1) - %param_2.4517 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.55 = s32[263]{0:T(512)} fusion(%param_0.4567, %param_1.5346, %param_2.4517), kind=kCustom, calls=%fused_computation.31.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.30 (param_0.4566: s32[263], param_1.5350: s32[8], param_2.4519: s32[8]) -> s32[263] { + %param_0.4566 = s32[263]{0:T(512)} parameter(0) + %param_1.5350 = s32[8]{0:T(128)} parameter(1) + %param_2.4519 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.55 = s32[263]{0:T(512)} fusion(%param_0.4566, %param_1.5350, %param_2.4519), kind=kCustom, calls=%fused_computation.31.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.30 (param_0.4568: s32[263], param_1.5347: s32[8], param_2.4518: s32[8]) -> s32[263] { - %param_0.4568 = s32[263]{0:T(512)} parameter(0) - %param_1.5347 = s32[8]{0:T(128)} parameter(1) - %param_2.4518 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.56.cloned.1 = s32[263]{0:T(512)} call(%param_0.4568, %param_1.5347, %param_2.4518), to_apply=%called_computation.30, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} +%async_computation.30 (param_0.4567: s32[263], param_1.5351: s32[8], param_2.4520: s32[8]) -> s32[263] { + %param_0.4567 = s32[263]{0:T(512)} parameter(0) + %param_1.5351 = s32[8]{0:T(128)} parameter(1) + %param_2.4520 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.56.cloned.1 = s32[263]{0:T(512)} call(%param_0.4567, %param_1.5351, %param_2.4520), to_apply=%called_computation.30, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.8 (param_0.108: s32[263], param_1.160: s32[8], param_2.104: s32[8], param_3.3107: token[]) -> s32[263] { - %param_3.3107 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.8 (param_0.108: s32[263], param_1.160: s32[8], param_2.104: s32[8], param_3.3109: token[]) -> s32[263] { + %param_3.3109 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.108 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.160 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.104 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1046,12 +1046,12 @@ StackFrames ROOT %scatter_offload_custom_fusion.56.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.56.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.8 (param_0.109: s32[263], param_1.161: s32[8], param_2.105: s32[8], param_3.3106: token[]) -> s32[263] { - %param_3.3106 = token[] parameter(3) +%async_computation.8 (param_0.109: s32[263], param_1.161: s32[8], param_2.105: s32[8], param_3.3108: token[]) -> s32[263] { + %param_3.3108 = token[] parameter(3) %param_0.109 = s32[263]{0:T(512)} parameter(0) %param_1.161 = s32[8]{0:T(128)} parameter(1) %param_2.105 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.26.cloned.1 = s32[263]{0:T(512)} call(%param_0.109, %param_1.161, %param_2.105, %param_3.3106), to_apply=%called_computation.8, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.26.cloned.1 = s32[263]{0:T(512)} call(%param_0.109, %param_1.161, %param_2.105, %param_3.3108), to_apply=%called_computation.8, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_14.20 (scatter-add.0: s32[], scatter-add.1: s32[]) -> s32[] { @@ -1060,33 +1060,33 @@ StackFrames ROOT %add.1312 = s32[]{:T(128)S(7)} add(%scatter-add.0, %scatter-add.1), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.17.clone.clone.clone (param_0.4571: s32[256], param_1.5348: s32[4096], param_2.4519: s32[4096]) -> s32[256] { - %param_0.4571 = s32[256]{0:T(256)} parameter(0) - %param_1.5348 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3941 = s32[4096]{0:T(1024)} reshape(%param_1.5348), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} +%fused_computation.17.clone.clone.clone (param_0.4570: s32[256], param_1.5352: s32[4096], param_2.4521: s32[4096]) -> s32[256] { + %param_0.4570 = s32[256]{0:T(256)} parameter(0) + %param_1.5352 = s32[4096]{0:T(1024)} parameter(1) + %reshape.3941 = s32[4096]{0:T(1024)} reshape(%param_1.5352), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} %transpose.1118 = s32[4096]{0:T(1024)} transpose(%reshape.3941), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} - %param_2.4519 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3942 = s32[4096]{0:T(1024)} reshape(%param_2.4519), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + %param_2.4521 = s32[4096]{0:T(1024)} parameter(2) + %reshape.3942 = s32[4096]{0:T(1024)} reshape(%param_2.4521), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} %transpose.1119 = s32[4096]{0:T(1024)} transpose(%reshape.3942), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.246 = s32[256]{0:T(256)} scatter(%param_0.4571, %transpose.1118, %transpose.1119), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_14.20, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter-add.246 = s32[256]{0:T(256)} scatter(%param_0.4570, %transpose.1118, %transpose.1119), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_14.20, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.32 (param_0.4572: s32[256], param_1.5349: s32[4096], param_2.4520: s32[4096]) -> s32[256] { - %param_0.4572 = s32[256]{0:T(256)} parameter(0) - %param_1.5349 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4520 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.57 = s32[256]{0:T(256)} fusion(%param_0.4572, %param_1.5349, %param_2.4520), kind=kCustom, calls=%fused_computation.17.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.32 (param_0.4571: s32[256], param_1.5353: s32[4096], param_2.4522: s32[4096]) -> s32[256] { + %param_0.4571 = s32[256]{0:T(256)} parameter(0) + %param_1.5353 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4522 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.57 = s32[256]{0:T(256)} fusion(%param_0.4571, %param_1.5353, %param_2.4522), kind=kCustom, calls=%fused_computation.17.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.32 (param_0.4573: s32[256], param_1.5350: s32[4096], param_2.4521: s32[4096]) -> s32[256] { - %param_0.4573 = s32[256]{0:T(256)} parameter(0) - %param_1.5350 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4521 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.58.cloned.1 = s32[256]{0:T(256)} call(%param_0.4573, %param_1.5350, %param_2.4521), to_apply=%called_computation.32, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} +%async_computation.32 (param_0.4572: s32[256], param_1.5354: s32[4096], param_2.4523: s32[4096]) -> s32[256] { + %param_0.4572 = s32[256]{0:T(256)} parameter(0) + %param_1.5354 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4523 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.58.cloned.1 = s32[256]{0:T(256)} call(%param_0.4572, %param_1.5354, %param_2.4523), to_apply=%called_computation.32, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.9 (param_0.111: s32[256], param_1.163: s32[4096], param_2.107: s32[4096], param_3.3085: token[]) -> s32[256] { - %param_3.3085 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.9 (param_0.111: s32[256], param_1.163: s32[4096], param_2.107: s32[4096], param_3.3087: token[]) -> s32[256] { + %param_3.3087 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.111 = s32[256]{0:T(256)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.163 = s32[4096]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.107 = s32[4096]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1094,22 +1094,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.58.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.58.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.9 (param_0.112: s32[256], param_1.164: s32[4096], param_2.108: s32[4096], param_3.3084: token[]) -> s32[256] { - %param_3.3084 = token[] parameter(3) +%async_computation.9 (param_0.112: s32[256], param_1.164: s32[4096], param_2.108: s32[4096], param_3.3086: token[]) -> s32[256] { + %param_3.3086 = token[] parameter(3) %param_0.112 = s32[256]{0:T(256)} parameter(0) %param_1.164 = s32[4096]{0:T(1024)} parameter(1) %param_2.108 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.29.cloned.1 = s32[256]{0:T(256)} call(%param_0.112, %param_1.164, %param_2.108, %param_3.3084), to_apply=%called_computation.9, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.29.cloned.1 = s32[256]{0:T(256)} call(%param_0.112, %param_1.164, %param_2.108, %param_3.3086), to_apply=%called_computation.9, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.33 (param_0.4574: s32[263]) -> s32[263] { - %param_0.4574 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2093 = s32[263]{0:T(512)} copy(%param_0.4574), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.33 (param_0.4573: s32[263]) -> s32[263] { + %param_0.4573 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2093 = s32[263]{0:T(512)} copy(%param_0.4573), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.33 (param_0.4575: s32[263]) -> s32[263] { - %param_0.4575 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2094.cloned.1 = s32[263]{0:T(512)} call(%param_0.4575), to_apply=%called_computation.33 +%async_computation.33 (param_0.4574: s32[263]) -> s32[263] { + %param_0.4574 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2094.cloned.1 = s32[263]{0:T(512)} call(%param_0.4574), to_apply=%called_computation.33 }, execution_thread="sparsecore" %region_20.26.clone.1 (scatter-add.141: s32[], scatter-add.142: s32[]) -> s32[] { @@ -1118,33 +1118,33 @@ StackFrames ROOT %add.2463 = s32[]{:T(128)S(7)} add(%scatter-add.141, %scatter-add.142), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.18.clone.clone.clone (param_0.4576: s32[263], param_1.5351: s32[256], param_2.4522: s32[256]) -> s32[263] { - %param_0.4576 = s32[263]{0:T(512)} parameter(0) - %param_1.5351 = s32[256]{0:T(256)} parameter(1) - %reshape.3943 = s32[256]{0:T(256)} reshape(%param_1.5351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} +%fused_computation.18.clone.clone.clone (param_0.4575: s32[263], param_1.5355: s32[256], param_2.4524: s32[256]) -> s32[263] { + %param_0.4575 = s32[263]{0:T(512)} parameter(0) + %param_1.5355 = s32[256]{0:T(256)} parameter(1) + %reshape.3943 = s32[256]{0:T(256)} reshape(%param_1.5355), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} %transpose.1120 = s32[256]{0:T(256)} transpose(%reshape.3943), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4522 = s32[256]{0:T(256)} parameter(2) - %reshape.3944 = s32[256]{0:T(256)} reshape(%param_2.4522) + %param_2.4524 = s32[256]{0:T(256)} parameter(2) + %reshape.3944 = s32[256]{0:T(256)} reshape(%param_2.4524) %transpose.1121 = s32[256]{0:T(256)} transpose(%reshape.3944), dimensions={0} - ROOT %scatter-add.247 = s32[263]{0:T(512)} scatter(%param_0.4576, %transpose.1120, %transpose.1121), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_20.26.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.247 = s32[263]{0:T(512)} scatter(%param_0.4575, %transpose.1120, %transpose.1121), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_20.26.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.34 (param_0.4577: s32[263], param_1.5352: s32[256], param_2.4523: s32[256]) -> s32[263] { - %param_0.4577 = s32[263]{0:T(512)} parameter(0) - %param_1.5352 = s32[256]{0:T(256)} parameter(1) - %param_2.4523 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.59 = s32[263]{0:T(512)} fusion(%param_0.4577, %param_1.5352, %param_2.4523), kind=kCustom, calls=%fused_computation.18.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.34 (param_0.4576: s32[263], param_1.5356: s32[256], param_2.4525: s32[256]) -> s32[263] { + %param_0.4576 = s32[263]{0:T(512)} parameter(0) + %param_1.5356 = s32[256]{0:T(256)} parameter(1) + %param_2.4525 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.59 = s32[263]{0:T(512)} fusion(%param_0.4576, %param_1.5356, %param_2.4525), kind=kCustom, calls=%fused_computation.18.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.34 (param_0.4578: s32[263], param_1.5353: s32[256], param_2.4524: s32[256]) -> s32[263] { - %param_0.4578 = s32[263]{0:T(512)} parameter(0) - %param_1.5353 = s32[256]{0:T(256)} parameter(1) - %param_2.4524 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.60.cloned.1 = s32[263]{0:T(512)} call(%param_0.4578, %param_1.5353, %param_2.4524), to_apply=%called_computation.34, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +%async_computation.34 (param_0.4577: s32[263], param_1.5357: s32[256], param_2.4526: s32[256]) -> s32[263] { + %param_0.4577 = s32[263]{0:T(512)} parameter(0) + %param_1.5357 = s32[256]{0:T(256)} parameter(1) + %param_2.4526 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.60.cloned.1 = s32[263]{0:T(512)} call(%param_0.4577, %param_1.5357, %param_2.4526), to_apply=%called_computation.34, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.10 (param_0.114: s32[263], param_1.166: s32[256], param_2.110: s32[256], param_3.3087: token[]) -> s32[263] { - %param_3.3087 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.10 (param_0.114: s32[263], param_1.166: s32[256], param_2.110: s32[256], param_3.3089: token[]) -> s32[263] { + %param_3.3089 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.114 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.166 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.110 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1154,22 +1154,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.60.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.60.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.10 (param_0.115: s32[263], param_1.167: s32[256], param_2.111: s32[256], param_3.3086: token[]) -> s32[263] { - %param_3.3086 = token[] parameter(3) +%async_computation.10 (param_0.115: s32[263], param_1.167: s32[256], param_2.111: s32[256], param_3.3088: token[]) -> s32[263] { + %param_3.3088 = token[] parameter(3) %param_0.115 = s32[263]{0:T(512)} parameter(0) %param_1.167 = s32[256]{0:T(256)} parameter(1) %param_2.111 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.32.cloned.1 = s32[263]{0:T(512)} call(%param_0.115, %param_1.167, %param_2.111, %param_3.3086), to_apply=%called_computation.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.32.cloned.1 = s32[263]{0:T(512)} call(%param_0.115, %param_1.167, %param_2.111, %param_3.3088), to_apply=%called_computation.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.35 (param_0.4579: f32[9]) -> f32[9] { - %param_0.4579 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2095 = f32[9]{0:T(128)} copy(%param_0.4579), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.35 (param_0.4578: f32[9]) -> f32[9] { + %param_0.4578 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2095 = f32[9]{0:T(128)} copy(%param_0.4578), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.35 (param_0.4580: f32[9]) -> f32[9] { - %param_0.4580 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2096.cloned.1 = f32[9]{0:T(128)} call(%param_0.4580), to_apply=%called_computation.35 +%async_computation.35 (param_0.4579: f32[9]) -> f32[9] { + %param_0.4579 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2096.cloned.1 = f32[9]{0:T(128)} call(%param_0.4579), to_apply=%called_computation.35 }, execution_thread="sparsecore" %region_26.33.clone.1 (scatter-add.145: f32[], scatter-add.146: f32[]) -> f32[] { @@ -1178,33 +1178,33 @@ StackFrames ROOT %add.2465 = f32[]{:T(128)S(7)} add(%scatter-add.145, %scatter-add.146), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.19.clone.clone.clone (param_0.4581: f32[9], param_1.5354: s32[256], param_2.4525: f32[256]) -> f32[9] { - %param_0.4581 = f32[9]{0:T(128)} parameter(0) - %param_1.5354 = s32[256]{0:T(256)} parameter(1) - %reshape.3945 = s32[256]{0:T(256)} reshape(%param_1.5354), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} +%fused_computation.19.clone.clone.clone (param_0.4580: f32[9], param_1.5358: s32[256], param_2.4527: f32[256]) -> f32[9] { + %param_0.4580 = f32[9]{0:T(128)} parameter(0) + %param_1.5358 = s32[256]{0:T(256)} parameter(1) + %reshape.3945 = s32[256]{0:T(256)} reshape(%param_1.5358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} %transpose.1122 = s32[256]{0:T(256)} transpose(%reshape.3945), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4525 = f32[256]{0:T(256)} parameter(2) - %reshape.3946 = f32[256]{0:T(256)} reshape(%param_2.4525) + %param_2.4527 = f32[256]{0:T(256)} parameter(2) + %reshape.3946 = f32[256]{0:T(256)} reshape(%param_2.4527) %transpose.1123 = f32[256]{0:T(256)} transpose(%reshape.3946), dimensions={0} - ROOT %scatter-add.248 = f32[9]{0:T(128)} scatter(%param_0.4581, %transpose.1122, %transpose.1123), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_26.33.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.248 = f32[9]{0:T(128)} scatter(%param_0.4580, %transpose.1122, %transpose.1123), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_26.33.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.36 (param_0.4582: f32[9], param_1.5355: s32[256], param_2.4526: f32[256]) -> f32[9] { - %param_0.4582 = f32[9]{0:T(128)} parameter(0) - %param_1.5355 = s32[256]{0:T(256)} parameter(1) - %param_2.4526 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.61 = f32[9]{0:T(128)} fusion(%param_0.4582, %param_1.5355, %param_2.4526), kind=kCustom, calls=%fused_computation.19.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.36 (param_0.4581: f32[9], param_1.5359: s32[256], param_2.4528: f32[256]) -> f32[9] { + %param_0.4581 = f32[9]{0:T(128)} parameter(0) + %param_1.5359 = s32[256]{0:T(256)} parameter(1) + %param_2.4528 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.61 = f32[9]{0:T(128)} fusion(%param_0.4581, %param_1.5359, %param_2.4528), kind=kCustom, calls=%fused_computation.19.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.36 (param_0.4583: f32[9], param_1.5356: s32[256], param_2.4527: f32[256]) -> f32[9] { - %param_0.4583 = f32[9]{0:T(128)} parameter(0) - %param_1.5356 = s32[256]{0:T(256)} parameter(1) - %param_2.4527 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.62.cloned.1 = f32[9]{0:T(128)} call(%param_0.4583, %param_1.5356, %param_2.4527), to_apply=%called_computation.36, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +%async_computation.36 (param_0.4582: f32[9], param_1.5360: s32[256], param_2.4529: f32[256]) -> f32[9] { + %param_0.4582 = f32[9]{0:T(128)} parameter(0) + %param_1.5360 = s32[256]{0:T(256)} parameter(1) + %param_2.4529 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.62.cloned.1 = f32[9]{0:T(128)} call(%param_0.4582, %param_1.5360, %param_2.4529), to_apply=%called_computation.36, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.11 (param_0.117: f32[9], param_1.169: s32[256], param_2.113: f32[256], param_3.3093: token[]) -> f32[9] { - %param_3.3093 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.11 (param_0.117: f32[9], param_1.169: s32[256], param_2.113: f32[256], param_3.3095: token[]) -> f32[9] { + %param_3.3095 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.117 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.169 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.113 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1214,22 +1214,22 @@ StackFrames ROOT %scatter_offload_custom_fusion.62.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.62.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.11 (param_0.118: f32[9], param_1.170: s32[256], param_2.114: f32[256], param_3.3092: token[]) -> f32[9] { - %param_3.3092 = token[] parameter(3) +%async_computation.11 (param_0.118: f32[9], param_1.170: s32[256], param_2.114: f32[256], param_3.3094: token[]) -> f32[9] { + %param_3.3094 = token[] parameter(3) %param_0.118 = f32[9]{0:T(128)} parameter(0) %param_1.170 = s32[256]{0:T(256)} parameter(1) %param_2.114 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.35.cloned.1 = f32[9]{0:T(128)} call(%param_0.118, %param_1.170, %param_2.114, %param_3.3092), to_apply=%called_computation.11, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.35.cloned.1 = f32[9]{0:T(128)} call(%param_0.118, %param_1.170, %param_2.114, %param_3.3094), to_apply=%called_computation.11, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.37 (param_0.4584: s32[263]) -> s32[263] { - %param_0.4584 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2097 = s32[263]{0:T(512)} copy(%param_0.4584), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.37 (param_0.4583: s32[263]) -> s32[263] { + %param_0.4583 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2097 = s32[263]{0:T(512)} copy(%param_0.4583), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.37 (param_0.4585: s32[263]) -> s32[263] { - %param_0.4585 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2098.cloned.1 = s32[263]{0:T(512)} call(%param_0.4585), to_apply=%called_computation.37 +%async_computation.37 (param_0.4584: s32[263]) -> s32[263] { + %param_0.4584 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2098.cloned.1 = s32[263]{0:T(512)} call(%param_0.4584), to_apply=%called_computation.37 }, execution_thread="sparsecore" %region_28.35.clone.1 (scatter-add.149: s32[], scatter-add.150: s32[]) -> s32[] { @@ -1238,33 +1238,33 @@ StackFrames ROOT %add.2467 = s32[]{:T(128)S(7)} add(%scatter-add.149, %scatter-add.150), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.20.clone.clone.clone (param_0.4586: s32[263], param_1.5357: s32[8], param_2.4528: s32[8]) -> s32[263] { - %param_0.4586 = s32[263]{0:T(512)} parameter(0) - %param_1.5357 = s32[8]{0:T(128)} parameter(1) - %reshape.3947 = s32[8]{0:T(128)} reshape(%param_1.5357), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} +%fused_computation.20.clone.clone.clone (param_0.4585: s32[263], param_1.5361: s32[8], param_2.4530: s32[8]) -> s32[263] { + %param_0.4585 = s32[263]{0:T(512)} parameter(0) + %param_1.5361 = s32[8]{0:T(128)} parameter(1) + %reshape.3947 = s32[8]{0:T(128)} reshape(%param_1.5361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} %transpose.1124 = s32[8]{0:T(128)} transpose(%reshape.3947), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4528 = s32[8]{0:T(128)} parameter(2) - %reshape.3948 = s32[8]{0:T(128)} reshape(%param_2.4528) + %param_2.4530 = s32[8]{0:T(128)} parameter(2) + %reshape.3948 = s32[8]{0:T(128)} reshape(%param_2.4530) %transpose.1125 = s32[8]{0:T(128)} transpose(%reshape.3948), dimensions={0} - ROOT %scatter-add.249 = s32[263]{0:T(512)} scatter(%param_0.4586, %transpose.1124, %transpose.1125), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_28.35.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter-add.249 = s32[263]{0:T(512)} scatter(%param_0.4585, %transpose.1124, %transpose.1125), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_28.35.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.38 (param_0.4587: s32[263], param_1.5358: s32[8], param_2.4529: s32[8]) -> s32[263] { - %param_0.4587 = s32[263]{0:T(512)} parameter(0) - %param_1.5358 = s32[8]{0:T(128)} parameter(1) - %param_2.4529 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.63 = s32[263]{0:T(512)} fusion(%param_0.4587, %param_1.5358, %param_2.4529), kind=kCustom, calls=%fused_computation.20.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.38 (param_0.4586: s32[263], param_1.5362: s32[8], param_2.4531: s32[8]) -> s32[263] { + %param_0.4586 = s32[263]{0:T(512)} parameter(0) + %param_1.5362 = s32[8]{0:T(128)} parameter(1) + %param_2.4531 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.63 = s32[263]{0:T(512)} fusion(%param_0.4586, %param_1.5362, %param_2.4531), kind=kCustom, calls=%fused_computation.20.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.38 (param_0.4588: s32[263], param_1.5359: s32[8], param_2.4530: s32[8]) -> s32[263] { - %param_0.4588 = s32[263]{0:T(512)} parameter(0) - %param_1.5359 = s32[8]{0:T(128)} parameter(1) - %param_2.4530 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.64.cloned.1 = s32[263]{0:T(512)} call(%param_0.4588, %param_1.5359, %param_2.4530), to_apply=%called_computation.38, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +%async_computation.38 (param_0.4587: s32[263], param_1.5363: s32[8], param_2.4532: s32[8]) -> s32[263] { + %param_0.4587 = s32[263]{0:T(512)} parameter(0) + %param_1.5363 = s32[8]{0:T(128)} parameter(1) + %param_2.4532 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.64.cloned.1 = s32[263]{0:T(512)} call(%param_0.4587, %param_1.5363, %param_2.4532), to_apply=%called_computation.38, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.12 (param_0.120: s32[263], param_1.172: s32[8], param_2.116: s32[8], param_3.3101: token[]) -> s32[263] { - %param_3.3101 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.12 (param_0.120: s32[263], param_1.172: s32[8], param_2.116: s32[8], param_3.3103: token[]) -> s32[263] { + %param_3.3103 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.120 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.172 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.116 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1274,12 +1274,12 @@ StackFrames ROOT %scatter_offload_custom_fusion.64.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.64.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.12 (param_0.121: s32[263], param_1.173: s32[8], param_2.117: s32[8], param_3.3100: token[]) -> s32[263] { - %param_3.3100 = token[] parameter(3) +%async_computation.12 (param_0.121: s32[263], param_1.173: s32[8], param_2.117: s32[8], param_3.3102: token[]) -> s32[263] { + %param_3.3102 = token[] parameter(3) %param_0.121 = s32[263]{0:T(512)} parameter(0) %param_1.173 = s32[8]{0:T(128)} parameter(1) %param_2.117 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.38.cloned.1 = s32[263]{0:T(512)} call(%param_0.121, %param_1.173, %param_2.117, %param_3.3100), to_apply=%called_computation.12, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.38.cloned.1 = s32[263]{0:T(512)} call(%param_0.121, %param_1.173, %param_2.117, %param_3.3102), to_apply=%called_computation.12, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_154.179 (reduce_sum.431: f32[], reduce_sum.254: f32[]) -> f32[] { @@ -1288,18 +1288,18 @@ StackFrames ROOT %reduce_sum.258 = f32[]{:T(128)} add(%reduce_sum.431, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.466 (param_0.4171: f32[3,1536,128,192]) -> f32[] { - %param_0.4171 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(0) - %bitcast.670 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4171), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3798 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.670, %bitcast.670), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.5104 = f32[]{:T(128)} constant(0) - ROOT %reduce.669 = f32[]{:T(128)} reduce(%mul.3798, %constant.5104), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.467 (param_0.4170: f32[3,1536,128,192]) -> f32[] { + %param_0.4170 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(0) + %bitcast.672 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4170), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %square.564 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.672, %bitcast.672), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5105 = f32[]{:T(128)} constant(0) + ROOT %reduce.669 = f32[]{:T(128)} reduce(%square.564, %constant.5105), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.467 (param_0.1419: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { - %param_0.1419 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) - %copy.1550 = bf16[1536,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1419), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\']"} - ROOT %bitcast.671 = bf16[3,1536,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1550), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} +%fused_computation.468 (param_0.1421: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { + %param_0.1421 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %copy.1550 = bf16[1536,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1421), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\']"} + ROOT %bitcast.673 = bf16[3,1536,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1550), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} } %region_221.246 (reduce_sum.893: f32[], reduce_sum.603: f32[]) -> f32[] { @@ -1314,54 +1314,54 @@ StackFrames ROOT %reduce_sum.450 = f32[]{:T(128)} add(%reduce_sum.655, %reduce_sum.449), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.468 (param_0.4141: f32[1536,3,128,192], param_1.5021: f32[], param_2.4296: f32[], param_3.2949: f32[], param_4.2203: f32[1536,3,128,192], param_5.2006: f32[], param_6.1443: f32[3,1536,128,192], param_7.1124: pred[], param_8.889: f32[1536,3,128,192]) -> (f32[], f32[1536,3,128,192], f32[1536,3,128,192], f32[1536,3,128,192], f32[]) { - %param_0.4141 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) - %param_3.2949 = f32[]{:T(128)S(6)} parameter(3) - %mul.5043.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2949), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.469 (param_0.4140: f32[1536,3,128,192], param_1.5025: f32[], param_2.4298: f32[], param_3.2951: f32[], param_4.2203: f32[1536,3,128,192], param_5.2006: f32[], param_6.1443: f32[3,1536,128,192], param_7.1124: pred[], param_8.889: f32[1536,3,128,192]) -> (f32[], f32[1536,3,128,192], f32[1536,3,128,192], f32[1536,3,128,192], f32[]) { + %param_0.4140 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %param_3.2951 = f32[]{:T(128)S(6)} parameter(3) + %mul.4727.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1124 = pred[]{:T(512)S(6)} parameter(7) %select_n.2165.clone.1 = pred[1536,3,128,192]{2,3,1,0:T(8,128)(4,1)} broadcast(%param_7.1124), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1443 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(6) - %bitcast.1372.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_6.1443), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1374.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_6.1443), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} %param_5.2006 = f32[]{:T(128)} parameter(5) %div.2575.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_5.2006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2574.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.1372.clone.1, %div.2575.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2164.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2165.clone.1, %bitcast.1372.clone.1, %div.2574.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4863.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4279.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4863.clone.1), dimensions={}, metadata={op_name="broadcast.334"} - %mul.5049.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2164.clone.1, %broadcast.4279.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2574.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.1374.clone.1, %div.2575.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2164.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2165.clone.1, %bitcast.1374.clone.1, %div.2574.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4864.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4279.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4864.clone.1), dimensions={}, metadata={op_name="broadcast.334"} + %mul.4733.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2164.clone.1, %broadcast.4279.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.889 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(8) - %constant.4867.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.5050.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4867.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5048.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.5050.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3443.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.5049.clone.1, %mul.5048.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4296 = f32[]{:T(128)S(6)} parameter(2) - %div.2571.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4296), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4868.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4734.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4868.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4732.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.4734.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3443.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4733.clone.1, %mul.4732.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4298 = f32[]{:T(128)S(6)} parameter(2) + %div.2571.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4298), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.399.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2164.clone.1, %select_n.2164.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4866.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.5047.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4866.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5045.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.5047.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4867.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4731.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4867.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4729.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.4731.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2203 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(4) - %constant.4865.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.5046.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4865.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5044.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.5046.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3442.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.5045.clone.1, %mul.5044.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5021 = f32[]{:T(128)S(6)} parameter(1) - %div.2570.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5021), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4866.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4730.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4866.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4728.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.4730.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3442.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4729.clone.1, %mul.4728.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5025 = f32[]{:T(128)S(6)} parameter(1) + %div.2570.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5025), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2569.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3442.clone.1, %div.2570.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.157.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} sqrt(%div.2569.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4864.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3441.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4864.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %constant.4865.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3441.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4865.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} %add.3440.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.157.clone.1, %add.3441.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1293.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%div.2571.clone.1, %add.3440.clone.1), metadata={op_name="multiply.290"} %div.2568.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3443.clone.1, %multiply.1293.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.5042.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4141, %broadcast.4279.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3439.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2568.clone.1, %mul.5042.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.5041.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.5043.clone.1, %add.3439.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3438.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4141, %mul.5041.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.330 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3438.clone.1, %add.3438.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5074 = f32[]{:T(128)} constant(0) - %reduce.670 = f32[]{:T(128)} reduce(%square.330, %constant.5074), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.671.clone.1 = f32[]{:T(128)} reduce(%integer_pow.399.clone.1, %constant.5074), dimensions={0,1,2,3}, to_apply=%region_187.212, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.4726.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4140, %broadcast.4279.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3439.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2568.clone.1, %mul.4726.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4725.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.4727.clone.1, %add.3439.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3438.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4140, %mul.4725.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.565 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3438.clone.1, %add.3438.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5075 = f32[]{:T(128)} constant(0) + %reduce.670 = f32[]{:T(128)} reduce(%square.565, %constant.5075), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.671.clone.1 = f32[]{:T(128)} reduce(%integer_pow.399.clone.1, %constant.5075), dimensions={0,1,2,3}, to_apply=%region_187.212, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.660 = (f32[]{:T(128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.670, %add.3438.clone.1, %add.3442.clone.1, %add.3443.clone.1, %reduce.671.clone.1) } @@ -1377,18 +1377,18 @@ StackFrames ROOT %reduce_sum.461 = f32[]{:T(128)} add(%reduce_sum.459, %reduce_sum.460), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.494 (param_0.4167: bf16[256,512,512], param_1.5043: bf16[256,512,512]) -> (f32[], f32[]) { - %param_0.4167 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) - %broadcast_in_dim.1358 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4167), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.693 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3827 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.693, %bitcast.693), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.5100 = f32[]{:T(128)} constant(0) - %reduce.672 = f32[]{:T(128)} reduce(%mul.3827, %constant.5100), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.5043 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(1) - %broadcast_in_dim.1366.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_1.5043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.701.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1366.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3833.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.701.clone.1, %bitcast.701.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.674.clone.1 = f32[]{:T(128)} reduce(%mul.3833.clone.1, %constant.5100), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.495 (param_0.4166: bf16[256,512,512], param_1.5047: bf16[256,512,512]) -> (f32[], f32[]) { + %param_0.4166 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %broadcast_in_dim.1358 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4166), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.695 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %square.570 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.695, %bitcast.695), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5101 = f32[]{:T(128)} constant(0) + %reduce.672 = f32[]{:T(128)} reduce(%square.570, %constant.5101), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.5047 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(1) + %broadcast_in_dim.1366.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_1.5047), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.703.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1366.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %square.576.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.703.clone.1, %bitcast.703.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.674.clone.1 = f32[]{:T(128)} reduce(%square.576.clone.1, %constant.5101), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.767 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.672, %reduce.674.clone.1) } @@ -1398,13 +1398,13 @@ StackFrames ROOT %reduce_sum.286 = f32[]{:T(128)} add(%reduce_sum.466, %reduce_sum.279), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.496 (param_0.4166: bf16[256,512,512]) -> f32[] { - %param_0.4166 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) - %broadcast_in_dim.1362 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4166), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.697 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3830 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.697, %bitcast.697), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.5099 = f32[]{:T(128)} constant(0) - ROOT %reduce.673 = f32[]{:T(128)} reduce(%mul.3830, %constant.5099), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.497 (param_0.4165: bf16[256,512,512]) -> f32[] { + %param_0.4165 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %broadcast_in_dim.1362 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4165), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.699 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %square.573 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.699, %bitcast.699), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5100 = f32[]{:T(128)} constant(0) + ROOT %reduce.673 = f32[]{:T(128)} reduce(%square.573, %constant.5100), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_227.252 (reduce_sum.935: f32[], reduce_sum.631: f32[]) -> f32[] { @@ -1419,61 +1419,61 @@ StackFrames ROOT %reduce_sum.472 = f32[]{:T(128)} add(%reduce_sum.697, %reduce_sum.471), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.514 (param_0.4135: f32[], param_1.5015: f32[256,1,512,512], param_2.4290: f32[], param_3.2943: f32[256,1,512,512], param_4.2197: f32[], param_5.2000: bf16[256,512,512], param_6.1437: pred[], param_7.1118: f32[], param_8.883: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.515 (param_0.4134: f32[], param_1.5019: f32[256,1,512,512], param_2.4292: f32[], param_3.2945: f32[256,1,512,512], param_4.2197: f32[], param_5.2000: bf16[256,512,512], param_6.1437: pred[], param_7.1118: f32[], param_8.883: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.883 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1357.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.883), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %bitcast.1359.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.883), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} %param_7.1118 = f32[]{:T(128)S(6)} parameter(7) - %mul.4992.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4676.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1437 = pred[]{:T(512)S(6)} parameter(6) %select_n.2147.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1437), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2000 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) %broadcast_in_dim.1572.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2000), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1359.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1572.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1361.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1572.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2197 = f32[]{:T(128)} parameter(4) %div.2533.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2197), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2532.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1359.clone.1, %div.2533.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2146.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2147.clone.1, %bitcast.1359.clone.1, %div.2532.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4833.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4259.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4833.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} - %mul.4994.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2146.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2943 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1358.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2943), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %constant.4832.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4258.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4832.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4993.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1358.clone.1, %broadcast.4258.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3408.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4994.clone.1, %mul.4993.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4290 = f32[]{:T(128)S(6)} parameter(2) - %div.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4290), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2532.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1361.clone.1, %div.2533.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2146.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2147.clone.1, %bitcast.1361.clone.1, %div.2532.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4834.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4259.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4834.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} + %mul.4678.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2146.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2945 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1360.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2945), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %constant.4833.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4258.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4833.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.4677.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4258.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3408.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4678.clone.1, %mul.4677.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4292 = f32[]{:T(128)S(6)} parameter(2) + %div.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4292), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2146.clone.1, %select_n.2146.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4831.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4261.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4831.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} - %mul.4996.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4261.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5015 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1360.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5015), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %constant.4830.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4260.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4830.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4995.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4260.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3409.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4996.clone.1, %mul.4995.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4135 = f32[]{:T(128)S(6)} parameter(0) - %div.2530.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4135), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4832.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4261.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4832.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} + %mul.4680.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4261.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5019 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1362.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5019), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %constant.4831.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4260.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4831.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.4679.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4260.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3409.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4680.clone.1, %mul.4679.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4134 = f32[]{:T(128)S(6)} parameter(0) + %div.2530.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4134), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2529.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3409.clone.1, %div.2530.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.151.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2529.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4834.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4257.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4834.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %constant.4835.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4257.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4835.clone.1), dimensions={}, metadata={op_name="broadcast.305"} %add.3407.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.151.clone.1, %broadcast.4257.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1287.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2531.clone.1, %add.3407.clone.1), metadata={op_name="multiply.296"} %div.2528.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3408.clone.1, %multiply.1287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4991.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1357.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3406.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2528.clone.1, %mul.4991.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4990.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4992.clone.1, %add.3406.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3405.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1357.clone.1, %mul.4990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.331 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3405.clone.1, %add.3405.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5068 = f32[]{:T(128)} constant(0) - %reduce.675 = f32[]{:T(128)} reduce(%square.331, %constant.5068), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.847.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3409.clone.1) - %bitcast.820.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3408.clone.1) - %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5068), dimensions={0,1,2,3}, to_apply=%region_193.218, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.670 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.675, %add.3405.clone.1, %bitcast.847.clone.1, %bitcast.820.clone.1, %reduce.684.clone.1) + %mul.4675.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1359.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3406.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2528.clone.1, %mul.4675.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4674.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4676.clone.1, %add.3406.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3405.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1359.clone.1, %mul.4674.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.577 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3405.clone.1, %add.3405.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5069 = f32[]{:T(128)} constant(0) + %reduce.675 = f32[]{:T(128)} reduce(%square.577, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.849.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3409.clone.1) + %bitcast.822.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3408.clone.1) + %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_193.218, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.670 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.675, %add.3405.clone.1, %bitcast.849.clone.1, %bitcast.822.clone.1, %reduce.684.clone.1) } %region_226.251 (reduce_sum.928: f32[], reduce_sum.625: f32[]) -> f32[] { @@ -1488,61 +1488,61 @@ StackFrames ROOT %reduce_sum.470 = f32[]{:T(128)} add(%reduce_sum.690, %reduce_sum.465), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.515 (param_0.4136: f32[], param_1.5016: f32[256,1,512,512], param_2.4291: f32[], param_3.2944: f32[256,1,512,512], param_4.2198: f32[], param_5.2001: bf16[256,512,512], param_6.1438: pred[], param_7.1119: f32[], param_8.884: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.516 (param_0.4135: f32[], param_1.5020: f32[256,1,512,512], param_2.4293: f32[], param_3.2946: f32[256,1,512,512], param_4.2198: f32[], param_5.2001: bf16[256,512,512], param_6.1438: pred[], param_7.1119: f32[], param_8.884: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.884 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1361.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.884), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %bitcast.1363.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.884), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} %param_7.1119 = f32[]{:T(128)S(6)} parameter(7) - %mul.4999.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4683.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1438 = pred[]{:T(512)S(6)} parameter(6) %select_n.2149.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1438), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2001 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) %broadcast_in_dim.1573.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2001), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1363.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1573.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1365.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1573.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2198 = f32[]{:T(128)} parameter(4) %div.2539.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2198), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2538.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1363.clone.1, %div.2539.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2148.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2149.clone.1, %bitcast.1363.clone.1, %div.2538.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4838.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4264.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4838.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} - %mul.5001.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2148.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2944 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1362.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2944), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %constant.4837.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4263.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4837.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.5000.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4263.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3413.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.5001.clone.1, %mul.5000.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4291 = f32[]{:T(128)S(6)} parameter(2) - %div.2537.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4291), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2538.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1365.clone.1, %div.2539.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2148.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2149.clone.1, %bitcast.1365.clone.1, %div.2538.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4839.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4264.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4839.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} + %mul.4685.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2148.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2946 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1364.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2946), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %constant.4838.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4263.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4838.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.4684.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4263.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3413.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4685.clone.1, %mul.4684.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4293 = f32[]{:T(128)S(6)} parameter(2) + %div.2537.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4293), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2148.clone.1, %select_n.2148.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4836.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4266.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4836.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} - %mul.5003.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4266.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5016 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1364.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5016), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %constant.4835.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4265.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4835.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.5002.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4265.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3414.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.5003.clone.1, %mul.5002.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4136 = f32[]{:T(128)S(6)} parameter(0) - %div.2536.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4136), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4837.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4266.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4837.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} + %mul.4687.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4266.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5020 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1366.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5020), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %constant.4836.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4265.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4836.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.4686.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4265.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3414.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4687.clone.1, %mul.4686.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4135 = f32[]{:T(128)S(6)} parameter(0) + %div.2536.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4135), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2535.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3414.clone.1, %div.2536.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.152.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2535.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4839.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4262.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4839.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %constant.4840.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4262.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="broadcast.305"} %add.3412.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.152.clone.1, %broadcast.4262.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1288.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2537.clone.1, %add.3412.clone.1), metadata={op_name="multiply.295"} %div.2534.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3413.clone.1, %multiply.1288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4998.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1361.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3411.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2534.clone.1, %mul.4998.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4997.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4999.clone.1, %add.3411.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3410.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1361.clone.1, %mul.4997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.332 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3410.clone.1, %add.3410.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5069 = f32[]{:T(128)} constant(0) - %reduce.676 = f32[]{:T(128)} reduce(%square.332, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.838.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3414.clone.1) - %bitcast.811.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3413.clone.1) - %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_192.217, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.669 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.676, %add.3410.clone.1, %bitcast.838.clone.1, %bitcast.811.clone.1, %reduce.685.clone.1) + %mul.4682.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1363.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3411.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2534.clone.1, %mul.4682.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4681.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4683.clone.1, %add.3411.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3410.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1363.clone.1, %mul.4681.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.578 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3410.clone.1, %add.3410.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5070 = f32[]{:T(128)} constant(0) + %reduce.676 = f32[]{:T(128)} reduce(%square.578, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.840.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3414.clone.1) + %bitcast.813.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3413.clone.1) + %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_192.217, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.669 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.676, %add.3410.clone.1, %bitcast.840.clone.1, %bitcast.813.clone.1, %reduce.685.clone.1) } %region_225.250 (reduce_sum.921: f32[], reduce_sum.619: f32[]) -> f32[] { @@ -1557,61 +1557,61 @@ StackFrames ROOT %reduce_sum.464 = f32[]{:T(128)} add(%reduce_sum.683, %reduce_sum.463), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.516 (param_0.4137: f32[], param_1.5017: f32[256,1,512,512], param_2.4292: f32[], param_3.2945: f32[256,1,512,512], param_4.2199: f32[], param_5.2002: bf16[256,512,512], param_6.1439: pred[], param_7.1120: f32[], param_8.885: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.517 (param_0.4136: f32[], param_1.5021: f32[256,1,512,512], param_2.4294: f32[], param_3.2947: f32[256,1,512,512], param_4.2199: f32[], param_5.2002: bf16[256,512,512], param_6.1439: pred[], param_7.1120: f32[], param_8.885: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.885 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1365.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.885), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %bitcast.1367.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.885), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} %param_7.1120 = f32[]{:T(128)S(6)} parameter(7) - %mul.5006.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4690.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1439 = pred[]{:T(512)S(6)} parameter(6) %select_n.2151.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1439), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2002 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) %broadcast_in_dim.1574.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2002), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1367.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1574.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1369.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1574.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2199 = f32[]{:T(128)} parameter(4) %div.2545.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2199), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2544.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1367.clone.1, %div.2545.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2150.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2151.clone.1, %bitcast.1367.clone.1, %div.2544.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4843.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4269.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4843.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} - %mul.5008.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2150.clone.1, %broadcast.4269.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2945 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1366.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2945), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4842.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4268.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4842.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.5007.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4268.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3418.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.5008.clone.1, %mul.5007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4292 = f32[]{:T(128)S(6)} parameter(2) - %div.2543.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4292), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2544.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1369.clone.1, %div.2545.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2150.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2151.clone.1, %bitcast.1369.clone.1, %div.2544.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4844.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4269.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4844.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} + %mul.4692.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2150.clone.1, %broadcast.4269.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2947 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1368.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2947), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %constant.4843.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4268.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4843.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.4691.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4268.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3418.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4692.clone.1, %mul.4691.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4294 = f32[]{:T(128)S(6)} parameter(2) + %div.2543.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4294), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2150.clone.1, %select_n.2150.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4841.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4271.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} - %mul.5010.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4271.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5017 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1368.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5017), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4840.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4270.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.5009.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4270.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3419.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.5010.clone.1, %mul.5009.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4137 = f32[]{:T(128)S(6)} parameter(0) - %div.2542.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4137), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4842.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4271.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4842.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} + %mul.4694.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4271.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5021 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1370.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5021), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %constant.4841.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4270.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.4693.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1370.clone.1, %broadcast.4270.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3419.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4694.clone.1, %mul.4693.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4136 = f32[]{:T(128)S(6)} parameter(0) + %div.2542.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4136), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2541.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3419.clone.1, %div.2542.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.153.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2541.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4844.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4267.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4844.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %constant.4845.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4267.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4845.clone.1), dimensions={}, metadata={op_name="broadcast.305"} %add.3417.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.153.clone.1, %broadcast.4267.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1289.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2543.clone.1, %add.3417.clone.1), metadata={op_name="multiply.294"} %div.2540.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3418.clone.1, %multiply.1289.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.5005.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1365.clone.1, %broadcast.4269.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3416.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2540.clone.1, %mul.5005.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.5004.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.5006.clone.1, %add.3416.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3415.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1365.clone.1, %mul.5004.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.333 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3415.clone.1, %add.3415.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5070 = f32[]{:T(128)} constant(0) - %reduce.677 = f32[]{:T(128)} reduce(%square.333, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.829.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3419.clone.1) - %bitcast.802.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3418.clone.1) - %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_191.216, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.668 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.677, %add.3415.clone.1, %bitcast.829.clone.1, %bitcast.802.clone.1, %reduce.686.clone.1) + %mul.4689.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1367.clone.1, %broadcast.4269.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3416.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2540.clone.1, %mul.4689.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4688.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4690.clone.1, %add.3416.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3415.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1367.clone.1, %mul.4688.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.579 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3415.clone.1, %add.3415.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5071 = f32[]{:T(128)} constant(0) + %reduce.677 = f32[]{:T(128)} reduce(%square.579, %constant.5071), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.831.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3419.clone.1) + %bitcast.804.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3418.clone.1) + %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5071), dimensions={0,1,2,3}, to_apply=%region_191.216, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.668 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.677, %add.3415.clone.1, %bitcast.831.clone.1, %bitcast.804.clone.1, %reduce.686.clone.1) } %region_155.180 (reduce_sum.438: f32[], reduce_sum.259: f32[]) -> f32[] { @@ -1620,61 +1620,61 @@ StackFrames ROOT %reduce_sum.260 = f32[]{:T(128)} add(%reduce_sum.438, %reduce_sum.259), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.528.clone.clone.clone (param_0.4080: bf16[4,128,129280], param_1.4949: s32[4,128], param_2.4223: f32[4,128], param_3.2911: f32[4,128], param_4.2170: bf16[4,128], param_5.1978: f32[4,128]) -> bf16[4,128,129280] { +%fused_computation.529.clone.clone.clone (param_0.4079: bf16[4,128,129280], param_1.4953: s32[4,128], param_2.4225: f32[4,128], param_3.2913: f32[4,128], param_4.2170: bf16[4,128], param_5.1978: f32[4,128]) -> bf16[4,128,129280] { %param_5.1978 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.5219 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.2911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.5218 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.4080 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.3161 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4080), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %mul.4903 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.2913 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.4902 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2913), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.4079 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.3163 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4079), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} %param_4.2170 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) %sub.804 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_4.2170), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.803 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.3161, %sub.804), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.803 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.3163, %sub.804), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.534 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.803), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.5217 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.5218, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.4223 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.2698 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4223), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.2697 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.5217, %div.2698), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.4949 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.371 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.4949), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.4901 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4902, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.4225 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.2698 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4225), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.2697 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.4901, %div.2698), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.4953 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.371 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.4953), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.370 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.369 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.371, %eq.370), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.3160 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%eq.369), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.802 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%div.2697, %convert_element_type.3160), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.5216 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.5219, %sub.802), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.3159 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.5216), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convert_element_type.3162 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%eq.369), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.802 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%div.2697, %convert_element_type.3162), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.4900 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4903, %sub.802), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.3161 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.4900), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} } -%fused_computation.938.clone.clone (param_0.4081: f32[4,128], param_1.4950: bf16[4,128,512], param_2.4225: bf16[512]) -> bf16[4,128,512] { - %param_2.4225 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) - %dot_general.831 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4225), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.4950 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3163 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.4081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.5221 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4081), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.5220 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3163, %mul.5221), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.3162 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.5220), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.830 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.831, %convert_element_type.3162), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.939.clone.clone (param_0.4080: f32[4,128], param_1.4954: bf16[4,128,512], param_2.4227: bf16[512]) -> bf16[4,128,512] { + %param_2.4227 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.831 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4227), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.4954 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.3165 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4954), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.4080 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.4905 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4080), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.4904 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3165, %mul.4905), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.3164 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4904), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.830 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.831, %convert_element_type.3164), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.517 (param_0.4170: bf16[4,128,129280], param_1.5045: s32[4,128], param_2.4317: f32[4,128], param_3.2967: f32[4,128], param_4.2219: bf16[4,128], param_5.2020: f32[4,128], param_6.1457: f32[4,128], param_7.1138: bf16[4,128,512], param_8.902: bf16[512]) -> (f32[], bf16[512,129280,1]) { +%fused_computation.518 (param_0.4169: bf16[4,128,129280], param_1.5049: s32[4,128], param_2.4319: f32[4,128], param_3.2969: f32[4,128], param_4.2219: bf16[4,128], param_5.2020: f32[4,128], param_6.1457: f32[4,128], param_7.1138: bf16[4,128,512], param_8.902: bf16[512]) -> (f32[], bf16[512,129280,1]) { %param_6.1457 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) %param_7.1138 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) %param_8.902 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(8) - %fusion.574.clone.1 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_6.1457, %param_7.1138, %param_8.902), kind=kLoop, calls=%fused_computation.938.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.4170 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.5045 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4317 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.2967 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %fusion.577.clone.1 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_6.1457, %param_7.1138, %param_8.902), kind=kLoop, calls=%fused_computation.939.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.4169 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.5049 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.4319 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.2969 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) %param_4.2219 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) %param_5.2020 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} fusion(%param_0.4170, %param_1.5045, %param_2.4317, %param_3.2967, %param_4.2219, /*index=5*/%param_5.2020), kind=kLoop, calls=%fused_computation.528.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.141.clone.1 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.574.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.774 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%convolution.141.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.2663 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.774), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %mul.3871 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2663, %convert_element_type.2663), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.5103 = f32[]{:T(128)} constant(0) - %reduce.678 = f32[]{:T(128)} reduce(%mul.3871, %constant.5103), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %multiply_convert_fusion.1.clone.1 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} fusion(%param_0.4169, %param_1.5049, %param_2.4319, %param_3.2969, %param_4.2219, /*index=5*/%param_5.2020), kind=kLoop, calls=%fused_computation.529.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convolution.141.clone.1 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.577.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %bitcast.776 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%convolution.141.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.2665 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.776), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %square.581 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2665, %convert_element_type.2665), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5104 = f32[]{:T(128)} constant(0) + %reduce.678 = f32[]{:T(128)} reduce(%square.581, %constant.5104), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.757 = (f32[]{:T(128)}, bf16[512,129280,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.678, %convolution.141.clone.1) } @@ -1684,12 +1684,12 @@ StackFrames ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.564, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.518 (param_0.4154: bf16[129280,512]) -> f32[] { - %param_0.4154 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2665 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4154), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %mul.3873 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2665, %convert_element_type.2665), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.5087 = f32[]{:T(128)} constant(0) - ROOT %reduce.679 = f32[]{:T(128)} reduce(%mul.3873, %constant.5087), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.519 (param_0.4153: bf16[129280,512]) -> f32[] { + %param_0.4153 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2667 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4153), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %square.583 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2667, %convert_element_type.2667), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5088 = f32[]{:T(128)} constant(0) + ROOT %reduce.679 = f32[]{:T(128)} reduce(%square.583, %constant.5088), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_240.265 (reduce_sum.1026: f32[], reduce_sum.689: f32[]) -> f32[] { @@ -1704,54 +1704,54 @@ StackFrames ROOT %reduce_sum.534 = f32[]{:T(128)} add(%reduce_sum.788, %reduce_sum.533), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.519 (param_0.4122: f32[129280,512], param_1.5002: f32[], param_2.4277: f32[], param_3.2930: f32[], param_4.2184: f32[129280,512], param_5.1987: f32[], param_6.1424: bf16[129280,512], param_7.1105: pred[], param_8.870: f32[129280,512]) -> (f32[], f32[129280,512], f32[129280,512], f32[129280,512], f32[]) { - %param_0.4122 = f32[129280,512]{1,0:T(8,128)} parameter(0) - %param_3.2930 = f32[]{:T(128)S(6)} parameter(3) - %mul.4880.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.520 (param_0.4121: f32[129280,512], param_1.5006: f32[], param_2.4279: f32[], param_3.2932: f32[], param_4.2184: f32[129280,512], param_5.1987: f32[], param_6.1424: bf16[129280,512], param_7.1105: pred[], param_8.870: f32[129280,512]) -> (f32[], f32[129280,512], f32[129280,512], f32[129280,512], f32[]) { + %param_0.4121 = f32[129280,512]{1,0:T(8,128)} parameter(0) + %param_3.2932 = f32[]{:T(128)S(6)} parameter(3) + %mul.4564.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2932), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1105 = pred[]{:T(512)S(6)} parameter(7) %select_n.2105.clone.1 = pred[129280,512]{1,0:T(8,128)(4,1)} broadcast(%param_7.1105), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1424 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.3104.clone.1 = f32[129280,512]{1,0:T(8,128)} convert(%param_6.1424), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convert_element_type.3106.clone.1 = f32[129280,512]{1,0:T(8,128)} convert(%param_6.1424), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %param_5.1987 = f32[]{:T(128)} parameter(5) %div.2439.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_5.1987), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2438.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%convert_element_type.3104.clone.1, %div.2439.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2104.clone.1 = f32[129280,512]{1,0:T(8,128)} select(%select_n.2105.clone.1, %convert_element_type.3104.clone.1, %div.2438.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4753.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4209.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4753.clone.1), dimensions={}, metadata={op_name="broadcast.318"} - %mul.4886.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4209.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2438.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%convert_element_type.3106.clone.1, %div.2439.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2104.clone.1 = f32[129280,512]{1,0:T(8,128)} select(%select_n.2105.clone.1, %convert_element_type.3106.clone.1, %div.2438.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4754.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4209.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4754.clone.1), dimensions={}, metadata={op_name="broadcast.318"} + %mul.4570.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4209.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.870 = f32[129280,512]{1,0:T(8,128)} parameter(8) - %constant.4757.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4887.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4757.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4885.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4887.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3338.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4886.clone.1, %mul.4885.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4277 = f32[]{:T(128)S(6)} parameter(2) - %div.2435.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4758.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4571.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4758.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4569.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3338.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4570.clone.1, %mul.4569.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4279 = f32[]{:T(128)S(6)} parameter(2) + %div.2435.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4279), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.380.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2104.clone.1, %select_n.2104.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4756.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4884.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4756.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4882.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4757.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4568.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4757.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4566.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4568.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2184 = f32[129280,512]{1,0:T(8,128)} parameter(4) - %constant.4755.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4883.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4755.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4881.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3337.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4882.clone.1, %mul.4881.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5002 = f32[]{:T(128)S(6)} parameter(1) - %div.2434.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5002), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4756.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4567.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4756.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4565.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4567.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3337.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4566.clone.1, %mul.4565.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5006 = f32[]{:T(128)S(6)} parameter(1) + %div.2434.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2433.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3337.clone.1, %div.2434.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.138.clone.1 = f32[129280,512]{1,0:T(8,128)} sqrt(%div.2433.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4754.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3336.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4754.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %constant.4755.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3336.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4755.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} %add.3335.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%sqrt.138.clone.1, %add.3336.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1274.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%div.2435.clone.1, %add.3335.clone.1), metadata={op_name="multiply.309"} %div.2432.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3338.clone.1, %multiply.1274.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4879.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4122, %broadcast.4209.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3334.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2432.clone.1, %mul.4879.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4878.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4880.clone.1, %add.3334.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3333.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4122, %mul.4878.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.334 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3333.clone.1, %add.3333.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5055 = f32[]{:T(128)} constant(0) - %reduce.680 = f32[]{:T(128)} reduce(%square.334, %constant.5055), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.687.clone.1 = f32[]{:T(128)} reduce(%integer_pow.380.clone.1, %constant.5055), dimensions={0,1}, to_apply=%region_206.231, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.4563.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4121, %broadcast.4209.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3334.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2432.clone.1, %mul.4563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4562.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4564.clone.1, %add.3334.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3333.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4121, %mul.4562.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.584 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3333.clone.1, %add.3333.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5056 = f32[]{:T(128)} constant(0) + %reduce.680 = f32[]{:T(128)} reduce(%square.584, %constant.5056), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.687.clone.1 = f32[]{:T(128)} reduce(%integer_pow.380.clone.1, %constant.5056), dimensions={0,1}, to_apply=%region_206.231, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.671 = (f32[]{:T(128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.680, %add.3333.clone.1, %add.3337.clone.1, %add.3338.clone.1, %reduce.687.clone.1) } @@ -1767,55 +1767,55 @@ StackFrames ROOT %reduce_sum.455 = f32[]{:T(128)} add(%reduce_sum.662, %reduce_sum.451), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.520 (param_0.4140: f32[512,129280], param_1.5020: f32[], param_2.4295: f32[], param_3.2948: f32[], param_4.2202: f32[512,129280], param_5.2005: f32[], param_6.1442: bf16[512,129280,1], param_7.1123: pred[], param_8.888: f32[512,129280]) -> (f32[], f32[512,129280], f32[512,129280], f32[512,129280], f32[]) { - %param_0.4140 = f32[512,129280]{1,0:T(8,128)} parameter(0) - %param_3.2948 = f32[]{:T(128)S(6)} parameter(3) - %mul.5033.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2948), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.521 (param_0.4139: f32[512,129280], param_1.5024: f32[], param_2.4297: f32[], param_3.2950: f32[], param_4.2202: f32[512,129280], param_5.2005: f32[], param_6.1442: bf16[512,129280,1], param_7.1123: pred[], param_8.888: f32[512,129280]) -> (f32[], f32[512,129280], f32[512,129280], f32[512,129280], f32[]) { + %param_0.4139 = f32[512,129280]{1,0:T(8,128)} parameter(0) + %param_3.2950 = f32[]{:T(128)S(6)} parameter(3) + %mul.4717.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2950), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1123 = pred[]{:T(512)S(6)} parameter(7) %select_n.2161.clone.1 = pred[512,129280]{1,0:T(8,128)(4,1)} broadcast(%param_7.1123), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1442 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.1370.clone.1 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%param_6.1442), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.3106.clone.1 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.1370.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %bitcast.1372.clone.1 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%param_6.1442), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.3108.clone.1 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.1372.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} %param_5.2005 = f32[]{:T(128)} parameter(5) %div.2567.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_5.2005), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2566.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%convert_element_type.3106.clone.1, %div.2567.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2160.clone.1 = f32[512,129280]{1,0:T(8,128)} select(%select_n.2161.clone.1, %convert_element_type.3106.clone.1, %div.2566.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4857.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4277.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4857.clone.1), dimensions={}, metadata={op_name="broadcast.333"} - %mul.5039.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2160.clone.1, %broadcast.4277.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2566.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%convert_element_type.3108.clone.1, %div.2567.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2160.clone.1 = f32[512,129280]{1,0:T(8,128)} select(%select_n.2161.clone.1, %convert_element_type.3108.clone.1, %div.2566.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4858.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4277.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="broadcast.333"} + %mul.4723.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2160.clone.1, %broadcast.4277.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.888 = f32[512,129280]{1,0:T(8,128)} parameter(8) - %constant.4861.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.5040.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4861.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5038.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.5040.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3437.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.5039.clone.1, %mul.5038.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4295 = f32[]{:T(128)S(6)} parameter(2) - %div.2563.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4295), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4862.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4724.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4862.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4722.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.4724.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3437.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4723.clone.1, %mul.4722.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4297 = f32[]{:T(128)S(6)} parameter(2) + %div.2563.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4297), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.398.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2160.clone.1, %select_n.2160.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4860.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.5037.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5035.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.5037.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4861.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4721.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4861.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4719.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.4721.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2202 = f32[512,129280]{1,0:T(8,128)} parameter(4) - %constant.4859.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.5036.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4859.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5034.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.5036.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3436.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.5035.clone.1, %mul.5034.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5020 = f32[]{:T(128)S(6)} parameter(1) - %div.2562.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5020), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4860.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4720.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4718.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.4720.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3436.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4719.clone.1, %mul.4718.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5024 = f32[]{:T(128)S(6)} parameter(1) + %div.2562.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5024), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2561.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3436.clone.1, %div.2562.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.156.clone.1 = f32[512,129280]{1,0:T(8,128)} sqrt(%div.2561.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4858.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3435.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %constant.4859.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3435.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4859.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} %add.3434.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%sqrt.156.clone.1, %add.3435.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1292.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%div.2563.clone.1, %add.3434.clone.1), metadata={op_name="multiply.291"} %div.2560.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3437.clone.1, %multiply.1292.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.5032.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4140, %broadcast.4277.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3433.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2560.clone.1, %mul.5032.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.5031.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.5033.clone.1, %add.3433.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3432.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4140, %mul.5031.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.335 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3432.clone.1, %add.3432.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5073 = f32[]{:T(128)} constant(0) - %reduce.681 = f32[]{:T(128)} reduce(%square.335, %constant.5073), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.688.clone.1 = f32[]{:T(128)} reduce(%integer_pow.398.clone.1, %constant.5073), dimensions={0,1}, to_apply=%region_188.213, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.4716.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4139, %broadcast.4277.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3433.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2560.clone.1, %mul.4716.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4715.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.4717.clone.1, %add.3433.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3432.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4139, %mul.4715.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.585 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3432.clone.1, %add.3432.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5074 = f32[]{:T(128)} constant(0) + %reduce.681 = f32[]{:T(128)} reduce(%square.585, %constant.5074), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.688.clone.1 = f32[]{:T(128)} reduce(%integer_pow.398.clone.1, %constant.5074), dimensions={0,1}, to_apply=%region_188.213, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.672 = (f32[]{:T(128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.681, %add.3432.clone.1, %add.3436.clone.1, %add.3437.clone.1, %reduce.688.clone.1) } @@ -1825,23 +1825,23 @@ StackFrames ROOT %reduce_sum.540 = f32[]{:T(128)} add(%reduce_sum.795, %reduce_sum.535), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.521 (param_0.4191: bf16[4,128,129280], param_1.5059: f32[4,128], param_2.4327: s32[4,128], param_3.2975: bf16[4,128]) -> f32[4,128] { - %param_2.4327 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.307 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4327), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.522 (param_0.4190: bf16[4,128,129280], param_1.5063: f32[4,128], param_2.4329: s32[4,128], param_3.2977: bf16[4,128]) -> f32[4,128] { + %param_2.4329 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.307 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4329), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.302 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.301 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.307, %eq.302), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.4191 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2670 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.2975 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.665 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2975), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.656 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2670, %sub.665), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.5059 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.663 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5059), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_0.4190 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2672 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4190), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_3.2977 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) + %sub.665 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2977), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.656 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2672, %sub.665), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.5063 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.663 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5063), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.652 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%sub.656, %sub.663), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.5127 = f32[]{:T(128)} constant(0) - %broadcast.3784 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5127), dimensions={}, metadata={op_name="broadcast.518"} - %mul.3874 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.301, %sub.652, %broadcast.3784), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3874, %constant.5127), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %constant.5128 = f32[]{:T(128)} constant(0) + %broadcast.3784 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5128), dimensions={}, metadata={op_name="broadcast.518"} + %mul.3624 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.301, %sub.652, %broadcast.3784), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3624, %constant.5128), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_37.47 (reduce_sum.76: f32[], reduce_sum.80: f32[]) -> f32[] { @@ -1850,15 +1850,15 @@ StackFrames ROOT %reduce_sum.83 = f32[]{:T(128)} add(%reduce_sum.76, %reduce_sum.80), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.532 (param_0.4192: bf16[4,128,129280], param_1.5060: bf16[4,128]) -> f32[4,128] { - %param_0.4192 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2676 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.5060 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.666 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5060), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.662 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2676, %sub.666), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.533 (param_0.4191: bf16[4,128,129280], param_1.5064: bf16[4,128]) -> f32[4,128] { + %param_0.4191 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2678 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.5064 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.666 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5064), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.662 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2678, %sub.666), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.448 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.662), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.5128 = f32[]{:T(128)} constant(0) - ROOT %reduce.683 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.448, %constant.5128), dimensions={2}, to_apply=%region_37.47, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %constant.5129 = f32[]{:T(128)} constant(0) + ROOT %reduce.683 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.448, %constant.5129), dimensions={2}, to_apply=%region_37.47, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_152.177 (reduce_sum.417: f32[], reduce_sum.244: f32[]) -> f32[] { @@ -1867,18 +1867,18 @@ StackFrames ROOT %reduce_sum.251 = f32[]{:T(128)} add(%reduce_sum.417, %reduce_sum.244), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.540 (param_0.4173: f32[3,512,128,256]) -> f32[] { - %param_0.4173 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.750 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4173), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3895 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.750, %bitcast.750), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.5106 = f32[]{:T(128)} constant(0) - ROOT %reduce.689 = f32[]{:T(128)} reduce(%mul.3895, %constant.5106), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.541 (param_0.4172: f32[3,512,128,256]) -> f32[] { + %param_0.4172 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.752 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4172), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %square.588 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.752, %bitcast.752), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5107 = f32[]{:T(128)} constant(0) + ROOT %reduce.689 = f32[]{:T(128)} reduce(%square.588, %constant.5107), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.541 (param_0.1600: f32[512,3,128,256]) -> bf16[3,512,128,256] { - %param_0.1600 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) - %copy.1551 = bf16[512,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1600), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\']"} - ROOT %bitcast.751 = bf16[3,512,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1551), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} +%fused_computation.542 (param_0.1602: f32[512,3,128,256]) -> bf16[3,512,128,256] { + %param_0.1602 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) + %copy.1551 = bf16[512,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1602), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\']"} + ROOT %bitcast.753 = bf16[3,512,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1551), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} } %region_219.244 (reduce_sum.879: f32[], reduce_sum.591: f32[]) -> f32[] { @@ -1893,54 +1893,54 @@ StackFrames ROOT %reduce_sum.442 = f32[]{:T(128)} add(%reduce_sum.641, %reduce_sum.437), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.542 (param_0.4143: f32[512,3,128,256], param_1.5023: f32[], param_2.4298: f32[], param_3.2951: f32[], param_4.2205: f32[512,3,128,256], param_5.2008: f32[], param_6.1445: f32[3,512,128,256], param_7.1126: pred[], param_8.891: f32[512,3,128,256]) -> (f32[], f32[512,3,128,256], f32[512,3,128,256], f32[512,3,128,256], f32[]) { - %param_0.4143 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) - %param_3.2951 = f32[]{:T(128)S(6)} parameter(3) - %mul.5063.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.543 (param_0.4142: f32[512,3,128,256], param_1.5027: f32[], param_2.4300: f32[], param_3.2953: f32[], param_4.2205: f32[512,3,128,256], param_5.2008: f32[], param_6.1445: f32[3,512,128,256], param_7.1126: pred[], param_8.891: f32[512,3,128,256]) -> (f32[], f32[512,3,128,256], f32[512,3,128,256], f32[512,3,128,256], f32[]) { + %param_0.4142 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) + %param_3.2953 = f32[]{:T(128)S(6)} parameter(3) + %mul.4747.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2953), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1126 = pred[]{:T(512)S(6)} parameter(7) %select_n.2173.clone.1 = pred[512,3,128,256]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.1126), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1445 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.1376.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_6.1445), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1378.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_6.1445), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} %param_5.2008 = f32[]{:T(128)} parameter(5) %div.2591.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_5.2008), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2590.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.1376.clone.1, %div.2591.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2172.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2173.clone.1, %bitcast.1376.clone.1, %div.2590.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4875.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4283.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4875.clone.1), dimensions={}, metadata={op_name="broadcast.336"} - %mul.5069.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2172.clone.1, %broadcast.4283.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2590.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.1378.clone.1, %div.2591.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2172.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2173.clone.1, %bitcast.1378.clone.1, %div.2590.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4876.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4283.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4876.clone.1), dimensions={}, metadata={op_name="broadcast.336"} + %mul.4753.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2172.clone.1, %broadcast.4283.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.891 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(8) - %constant.4879.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.5070.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4879.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5068.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.5070.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3455.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.5069.clone.1, %mul.5068.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4298 = f32[]{:T(128)S(6)} parameter(2) - %div.2587.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4298), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4880.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4754.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4880.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4752.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.4754.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3455.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4753.clone.1, %mul.4752.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4300 = f32[]{:T(128)S(6)} parameter(2) + %div.2587.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4300), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.401.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2172.clone.1, %select_n.2172.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4878.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.5067.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4878.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5065.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.5067.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4879.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4751.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4879.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4749.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.4751.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2205 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(4) - %constant.4877.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.5066.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4877.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5064.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.5066.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3454.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.5065.clone.1, %mul.5064.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5023 = f32[]{:T(128)S(6)} parameter(1) - %div.2586.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5023), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.4878.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4750.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4878.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4748.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.4750.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3454.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4749.clone.1, %mul.4748.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5027 = f32[]{:T(128)S(6)} parameter(1) + %div.2586.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5027), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2585.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3454.clone.1, %div.2586.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.159.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} sqrt(%div.2585.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4876.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3453.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4876.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %constant.4877.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3453.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4877.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} %add.3452.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3453.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1295.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%div.2587.clone.1, %add.3452.clone.1), metadata={op_name="multiply.288"} %div.2584.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3455.clone.1, %multiply.1295.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.5062.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4143, %broadcast.4283.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3451.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2584.clone.1, %mul.5062.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.5061.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.5063.clone.1, %add.3451.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3450.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4143, %mul.5061.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.336 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3450.clone.1, %add.3450.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5076 = f32[]{:T(128)} constant(0) - %reduce.690 = f32[]{:T(128)} reduce(%square.336, %constant.5076), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.691.clone.1 = f32[]{:T(128)} reduce(%integer_pow.401.clone.1, %constant.5076), dimensions={0,1,2,3}, to_apply=%region_185.210, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.4746.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4142, %broadcast.4283.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3451.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2584.clone.1, %mul.4746.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4745.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.4747.clone.1, %add.3451.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3450.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4142, %mul.4745.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.589 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3450.clone.1, %add.3450.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5077 = f32[]{:T(128)} constant(0) + %reduce.690 = f32[]{:T(128)} reduce(%square.589, %constant.5077), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.691.clone.1 = f32[]{:T(128)} reduce(%integer_pow.401.clone.1, %constant.5077), dimensions={0,1,2,3}, to_apply=%region_185.210, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.667 = (f32[]{:T(128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.690, %add.3450.clone.1, %add.3454.clone.1, %add.3455.clone.1, %reduce.691.clone.1) } @@ -1950,39 +1950,39 @@ StackFrames ROOT %reduce_sum.386 = f32[]{:T(128)} add(%reduce_sum.557, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.782.clone.clone (param_0.4107: f32[4,128], param_1.4994: bf16[4,128,1536], param_2.4259: bf16[1536]) -> bf16[4,128,1536,1] { - %param_2.4259 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.851 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4259), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.4994 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3185 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4994), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} - %param_0.4107 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.5267 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4107), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %mul.5266 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3185, %mul.5267), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %convert_element_type.3184 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.5266), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} - %dot_general.850 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.851, %convert_element_type.3184), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.1464 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dot_general.850), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.783.clone.clone (param_0.4106: f32[4,128], param_1.4998: bf16[4,128,1536], param_2.4261: bf16[1536]) -> bf16[4,128,1536,1] { + %param_2.4261 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.851 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4261), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.4998 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.3187 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4998), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} + %param_0.4106 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.4951 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4106), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} + %mul.4950 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3187, %mul.4951), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} + %convert_element_type.3186 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.4950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} + %dot_general.850 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.851, %convert_element_type.3186), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.1466 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dot_general.850), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} } %bitcast_fusion.12 (bitcast_input.12: bf16[4,128,128,192]) -> bf16[4,128,128,192] { %bitcast_input.12 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.1486 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} bitcast(%bitcast_input.12) -} - -%fused_computation.551 (param_0.4155: bf16[4,128,128,192], param_1.5034: f32[4,128], param_2.4309: bf16[4,128,1536], param_3.2962: bf16[1536]) -> (f32[], bf16[1536,128,192,1]) { - %param_1.5034 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4309 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.2962 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.457.clone.1 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.5034, %param_2.4309, %param_3.2962), kind=kLoop, calls=%fused_computation.782.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.4155 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - %fusion.748 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.4155), kind=kLoop, calls=%bitcast_fusion.12 - %convolution.144.clone.1 = bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)} convolution(%fusion.457.clone.1, %fusion.748), window={size=192x4 pad=191_191x0_0 rhs_reversal=1x0}, dim_labels=1fb0_1io0->bf01, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} - %bitcast.859 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%convolution.144.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} - %broadcast_in_dim.1388 = f32[1536,128,192]{1,0,2:T(8,128)} convert(%bitcast.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.761 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3904 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.761, %bitcast.761), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.5088 = f32[]{:T(128)} constant(0) - %reduce.692 = f32[]{:T(128)} reduce(%mul.3904, %constant.5088), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.766 = (f32[]{:T(128)}, bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)}) tuple(%reduce.692, %convolution.144.clone.1) + ROOT %bitcast.1488 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} bitcast(%bitcast_input.12) +} + +%fused_computation.552 (param_0.4154: bf16[4,128,128,192], param_1.5038: f32[4,128], param_2.4311: bf16[4,128,1536], param_3.2964: bf16[1536]) -> (f32[], bf16[1536,128,192,1]) { + %param_1.5038 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.4311 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.2964 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.460.clone.1 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.5038, %param_2.4311, %param_3.2964), kind=kLoop, calls=%fused_computation.783.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.4154 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) + %fusion.751 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.4154), kind=kLoop, calls=%bitcast_fusion.12 + %convolution.146.clone.1 = bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)} convolution(%fusion.460.clone.1, %fusion.751), window={size=192x4 pad=191_191x0_0 rhs_reversal=1x0}, dim_labels=1fb0_1io0->bf01, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} + %bitcast.861 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%convolution.146.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} + %broadcast_in_dim.1388 = f32[1536,128,192]{1,0,2:T(8,128)} convert(%bitcast.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.763 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %square.592 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.763, %bitcast.763), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5089 = f32[]{:T(128)} constant(0) + %reduce.692 = f32[]{:T(128)} reduce(%square.592, %constant.5089), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.766 = (f32[]{:T(128)}, bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)}) tuple(%reduce.692, %convolution.146.clone.1) } %region_239.264 (reduce_sum.1019: f32[], reduce_sum.687: f32[]) -> f32[] { @@ -1997,4 +1997,4 @@ StackFrames ROOT %reduce_sum.528 = f32[]{:T(128)} add(%reduce_sum.781, %reduce_sum.527), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.556 (param_0.4123: f32[], param_1.5003: f32[], param_2.4278: f32[], param_3.2931: f32[1536,1,128,192], param_4.2185: f32[1536,1,128,192], param_5.1988: f32[], param_6.1425: bf16[1536,128,192,1], param_7.1106: pred[], param_8.871: f32[1536,1,128,192]) -> (f32[], f32[1536,1,128,192], f32[1536,1,128,192], f32[1536,1,128,192], f32[]) { +%fused_computation.557 (param_0.4122: f32[], param_1.5007: f32[], param_2.4280: f32[], param_3.2933: f32[1536,1,128,192], param_4.2185: f32[1536,1,128,192], param_5.1988: f32[], param_6.1425: bf16[1536,128,192,1], param_7.1106: pred[], param_8.871: f32[1536,1,128,192]) -> (f32[], f32[1536,1,128,192], f32[1536,1,128,192], f32[1536,1,128,192], f32[]) { diff --git a/tests/utils/reference_hlo_llama3_8b.txt b/tests/utils/reference_hlo_llama3_8b.txt index 2a4c292f01..27c6529df2 100644 --- a/tests/utils/reference_hlo_llama3_8b.txt +++ b/tests/utils/reference_hlo_llama3_8b.txt @@ -44,62 +44,62 @@ StackFrames ROOT %reduce_sum.192 = f32[]{:T(128)} add(%reduce_sum.190, %reduce_sum.191), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.281.clone.clone.clone (param_0.1085: bf16[4,128,128256], param_1.1251: s32[4,128], param_2.1077: f32[4,128], param_3.781: f32[4,128], param_4.482: bf16[4,128], param_5.404: f32[4,128]) -> bf16[4,128,128256] { - %param_5.404 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1679 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.404), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.781 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1678 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.781), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1085 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1032 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1085), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.482 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.482), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1032, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.280.clone.clone.clone (param_0.1099: bf16[4,128,128256], param_1.1265: s32[4,128], param_2.1086: f32[4,128], param_3.785: f32[4,128], param_4.487: bf16[4,128], param_5.412: f32[4,128]) -> bf16[4,128,128256] { + %param_5.412 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.1613 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.785 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.1612 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.785), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1099 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1044 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1099), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.487 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.487), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1044, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.62 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1677 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1678, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1077 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1077), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1677, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1251 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1251), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.1611 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1612, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1611, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1265 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1265), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.48 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.47 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1031 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1031), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1676 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1679, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1030 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1676), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.317.clone.clone (param_0.1086: f32[4,128], param_1.1252: bf16[4,128,4096], param_2.1079: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1079 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1079), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1034 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1252), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1681 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1680 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1034, %mul.1681), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1033 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1680), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.220 (param_0.1105: bf16[4,128,128256], param_1.1267: s32[4,128], param_2.1103: f32[4,128], param_3.797: f32[4,128], param_4.497: bf16[4,128], param_5.419: f32[4,128], param_6.287: f32[4,128], param_7.186: bf16[4,128,4096], param_8.112: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { - %param_6.287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.186 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) - %param_8.112 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) - %fusion.229.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.287, %param_7.186, %param_8.112), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1105 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1267 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1103 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.797 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %param_4.497 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %param_5.419 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1105, %param_1.1267, %param_2.1103, %param_3.797, %param_4.497, /*index=5*/%param_5.419), kind=kLoop, calls=%fused_computation.281.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.82.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.229.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.300 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.82.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.911 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.300), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %mul.1350 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.911, %convert_element_type.911), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %convert_element_type.1043 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1043), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1610 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1613, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1042 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1610), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.316.clone.clone (param_0.1100: f32[4,128], param_1.1266: bf16[4,128,4096], param_2.1088: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1088 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.387 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1088), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1266 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1046 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1266), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1100 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1615 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1100), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1614 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1046, %mul.1615), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1045 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1614), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.387, %convert_element_type.1045), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.219 (param_0.1119: bf16[4,128,128256], param_1.1281: s32[4,128], param_2.1112: f32[4,128], param_3.801: f32[4,128], param_4.502: bf16[4,128], param_5.427: f32[4,128], param_6.299: f32[4,128], param_7.198: bf16[4,128,4096], param_8.116: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { + %param_6.299 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.198 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) + %param_8.116 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) + %fusion.239.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.299, %param_7.198, %param_8.116), kind=kLoop, calls=%fused_computation.316.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1119 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1281 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1112 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.801 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %param_4.502 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %param_5.427 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1119, %param_1.1281, %param_2.1112, %param_3.801, %param_4.502, /*index=5*/%param_5.427), kind=kLoop, calls=%fused_computation.280.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convolution.88.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.239.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %bitcast.306 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.88.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.923 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.306), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %square.157 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.923, %convert_element_type.923), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1006 = f32[]{:T(128)} constant(0) - %reduce.118 = f32[]{:T(128)} reduce(%mul.1350, %constant.1006), dimensions={0,1}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.82.clone.1) + %reduce.118 = f32[]{:T(128)} reduce(%square.157, %constant.1006), dimensions={0,1}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.88.clone.1) } %region_34.39 (reduce_sum.196: f32[], reduce_sum.197: f32[]) -> f32[] { @@ -108,12 +108,12 @@ StackFrames ROOT %reduce_sum.198 = f32[]{:T(128)} add(%reduce_sum.196, %reduce_sum.197), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.221 (param_0.1104: bf16[128256,4096]) -> f32[] { - %param_0.1104 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.913 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1104), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %mul.1352 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.913, %convert_element_type.913), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.220 (param_0.1118: bf16[128256,4096]) -> f32[] { + %param_0.1118 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.925 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1118), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %square.159 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.925, %convert_element_type.925), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1005 = f32[]{:T(128)} constant(0) - ROOT %reduce.119 = f32[]{:T(128)} reduce(%mul.1352, %constant.1005), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.119 = f32[]{:T(128)} reduce(%square.159, %constant.1005), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_60.65 (reduce_sum.338: f32[], reduce_sum.339: f32[]) -> f32[] { @@ -128,39 +128,39 @@ StackFrames ROOT %reduce_sum.261 = f32[]{:T(128)} add(%reduce_sum.259, %reduce_sum.260), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.222 (param_0.1092: f32[128256,4096], param_1.1255: f32[], param_2.1091: f32[], param_3.785: f32[], param_4.485: f32[128256,4096], param_5.407: f32[], param_6.275: bf16[128256,4096], param_7.174: pred[], param_8.100: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { - %param_0.1092 = f32[128256,4096]{1,0:T(8,128)} parameter(0) - %param_3.785 = f32[]{:T(128)S(6)} parameter(3) - %mul.1548.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.785), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.174 = pred[]{:T(512)S(6)} parameter(7) - %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.174), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.275 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1005.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.275), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %param_5.407 = f32[]{:T(128)} parameter(5) - %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.407), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1005.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1005.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.221 (param_0.1106: f32[128256,4096], param_1.1269: f32[], param_2.1100: f32[], param_3.789: f32[], param_4.490: f32[128256,4096], param_5.415: f32[], param_6.287: bf16[128256,4096], param_7.186: pred[], param_8.104: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { + %param_0.1106 = f32[128256,4096]{1,0:T(8,128)} parameter(0) + %param_3.789 = f32[]{:T(128)S(6)} parameter(3) + %mul.1482.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.186 = pred[]{:T(512)S(6)} parameter(7) + %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.186), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.287 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1017.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.287), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_5.415 = f32[]{:T(128)} parameter(5) + %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1017.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1017.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.907.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.554.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.907.clone.1), dimensions={}, metadata={op_name="broadcast.61"} - %mul.1554.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.100 = f32[128256,4096]{1,0:T(8,128)} parameter(8) + %mul.1488.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.104 = f32[128256,4096]{1,0:T(8,128)} parameter(8) %constant.911.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1555.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1553.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.100, %mul.1555.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1554.clone.1, %mul.1553.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1091 = f32[]{:T(128)S(6)} parameter(2) - %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1091), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1489.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1487.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.104, %mul.1489.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1488.clone.1, %mul.1487.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) + %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.60.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %select_n.241.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.910.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1552.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1550.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1552.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.485 = f32[128256,4096]{1,0:T(8,128)} parameter(4) + %mul.1486.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1484.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1486.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.490 = f32[128256,4096]{1,0:T(8,128)} parameter(4) %constant.909.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1551.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1549.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.485, %mul.1551.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1550.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1255 = f32[]{:T(128)S(6)} parameter(1) - %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1255), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1485.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1483.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.490, %mul.1485.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1484.clone.1, %mul.1483.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1269 = f32[]{:T(128)S(6)} parameter(1) + %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1269), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.719.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.775.clone.1, %div.720.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.58.clone.1 = f32[128256,4096]{1,0:T(8,128)} sqrt(%div.719.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.908.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -168,13 +168,13 @@ StackFrames %add.773.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%sqrt.58.clone.1, %add.774.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.256.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%div.721.clone.1, %add.773.clone.1), metadata={op_name="multiply.42"} %div.718.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.776.clone.1, %multiply.256.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1547.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1092, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1547.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1546.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1548.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1092, %mul.1546.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.118 = f32[128256,4096]{1,0:T(8,128)} multiply(%add.771.clone.1, %add.771.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1481.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1106, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1481.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1480.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1482.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1106, %mul.1480.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.160 = f32[128256,4096]{1,0:T(8,128)} multiply(%add.771.clone.1, %add.771.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.993 = f32[]{:T(128)} constant(0) - %reduce.120 = f32[]{:T(128)} reduce(%square.118, %constant.993), dimensions={0,1}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.120 = f32[]{:T(128)} reduce(%square.160, %constant.993), dimensions={0,1}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.122.clone.1 = f32[]{:T(128)} reduce(%integer_pow.60.clone.1, %constant.993), dimensions={0,1}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.135 = (f32[]{:T(128)}, f32[128256,4096]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.120, %add.771.clone.1, %add.775.clone.1, %add.776.clone.1, %reduce.122.clone.1) } @@ -191,40 +191,40 @@ StackFrames ROOT %reduce_sum.255 = f32[]{:T(128)} add(%reduce_sum.253, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.223 (param_0.1093: f32[4096,128256], param_1.1256: f32[], param_2.1092: f32[], param_3.786: f32[], param_4.486: f32[4096,128256], param_5.408: f32[], param_6.276: bf16[4096,128256,1], param_7.175: pred[], param_8.101: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { - %param_0.1093 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %param_3.786 = f32[]{:T(128)S(6)} parameter(3) - %mul.1558.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.786), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.175 = pred[]{:T(512)S(6)} parameter(7) - %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.175), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.276 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.403.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.1007.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.403.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %param_5.408 = f32[]{:T(128)} parameter(5) - %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.408), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1007.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1007.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.222 (param_0.1107: f32[4096,128256], param_1.1270: f32[], param_2.1101: f32[], param_3.790: f32[], param_4.491: f32[4096,128256], param_5.416: f32[], param_6.288: bf16[4096,128256,1], param_7.187: pred[], param_8.105: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { + %param_0.1107 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %param_3.790 = f32[]{:T(128)S(6)} parameter(3) + %mul.1492.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.187 = pred[]{:T(512)S(6)} parameter(7) + %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.187), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.288 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) + %bitcast.409.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.288), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.1019.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.409.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %param_5.416 = f32[]{:T(128)} parameter(5) + %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1019.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1019.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.913.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.556.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.913.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %mul.1564.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.101 = f32[4096,128256]{1,0:T(8,128)} parameter(8) + %mul.1498.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.105 = f32[4096,128256]{1,0:T(8,128)} parameter(8) %constant.917.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1565.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1563.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.101, %mul.1565.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1564.clone.1, %mul.1563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1092 = f32[]{:T(128)S(6)} parameter(2) - %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1092), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1499.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1497.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.105, %mul.1499.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1498.clone.1, %mul.1497.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) + %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.61.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %select_n.245.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.916.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1562.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1560.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.486 = f32[4096,128256]{1,0:T(8,128)} parameter(4) + %mul.1496.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1494.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1496.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.491 = f32[4096,128256]{1,0:T(8,128)} parameter(4) %constant.915.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1561.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1559.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.486, %mul.1561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1560.clone.1, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1256 = f32[]{:T(128)S(6)} parameter(1) - %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1256), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1495.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1493.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.491, %mul.1495.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1494.clone.1, %mul.1493.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1270 = f32[]{:T(128)S(6)} parameter(1) + %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1270), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.727.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.781.clone.1, %div.728.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.59.clone.1 = f32[4096,128256]{1,0:T(8,128)} sqrt(%div.727.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.914.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -232,13 +232,13 @@ StackFrames %add.779.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%sqrt.59.clone.1, %add.780.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.257.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%div.729.clone.1, %add.779.clone.1), metadata={op_name="multiply.41"} %div.726.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.782.clone.1, %multiply.257.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1557.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1093, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1556.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1558.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1093, %mul.1556.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.119 = f32[4096,128256]{1,0:T(8,128)} multiply(%add.777.clone.1, %add.777.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1491.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1107, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1491.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1490.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1492.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1107, %mul.1490.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.161 = f32[4096,128256]{1,0:T(8,128)} multiply(%add.777.clone.1, %add.777.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.994 = f32[]{:T(128)} constant(0) - %reduce.121 = f32[]{:T(128)} reduce(%square.119, %constant.994), dimensions={0,1}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.121 = f32[]{:T(128)} reduce(%square.161, %constant.994), dimensions={0,1}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.123.clone.1 = f32[]{:T(128)} reduce(%integer_pow.61.clone.1, %constant.994), dimensions={0,1}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.136 = (f32[]{:T(128)}, f32[4096,128256]{1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.121, %add.777.clone.1, %add.781.clone.1, %add.782.clone.1, %reduce.123.clone.1) } @@ -249,12 +249,12 @@ StackFrames ROOT %reduce_sum.156 = f32[]{:T(128)} add(%reduce_sum.154, %reduce_sum.155), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.240 (param_0.1110: f32[4,14336,4096]) -> f32[] { - %param_0.1110 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) - %bitcast.308 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1375 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.308, %bitcast.308), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.239 (param_0.1124: f32[4,14336,4096]) -> f32[] { + %param_0.1124 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) + %bitcast.314 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1124), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.164 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.314, %bitcast.314), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1011 = f32[]{:T(128)} constant(0) - ROOT %reduce.124 = f32[]{:T(128)} reduce(%mul.1375, %constant.1011), dimensions={0,1,2}, to_apply=%region_25.30, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.124 = f32[]{:T(128)} reduce(%square.164, %constant.1011), dimensions={0,1,2}, to_apply=%region_25.30, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_24.29 (reduce_sum.148: f32[], reduce_sum.149: f32[]) -> f32[] { @@ -269,35 +269,35 @@ StackFrames ROOT %reduce_sum.147 = f32[]{:T(128)} add(%reduce_sum.142, %reduce_sum.143), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.242 (param_0.1111: f32[4,4096,14336], param_1.1270: f32[4,4096,14336]) -> (f32[], f32[]) { - %param_0.1111 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) - %bitcast.312 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1378 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.312, %bitcast.312), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.241 (param_0.1125: f32[4,4096,14336], param_1.1284: f32[4,4096,14336]) -> (f32[], f32[]) { + %param_0.1125 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) + %bitcast.318 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1125), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.167 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.318, %bitcast.318), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1012 = f32[]{:T(128)} constant(0) - %reduce.125 = f32[]{:T(128)} reduce(%mul.1378, %constant.1012), dimensions={0,1,2}, to_apply=%region_24.29, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1270 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) - %bitcast.316.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1270), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1381.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.316.clone.1, %bitcast.316.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.126.clone.1 = f32[]{:T(128)} reduce(%mul.1381.clone.1, %constant.1012), dimensions={0,1,2}, to_apply=%region_23.28, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.125 = f32[]{:T(128)} reduce(%square.167, %constant.1012), dimensions={0,1,2}, to_apply=%region_24.29, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) + %bitcast.322.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.170.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.322.clone.1, %bitcast.322.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.126.clone.1 = f32[]{:T(128)} reduce(%square.170.clone.1, %constant.1012), dimensions={0,1,2}, to_apply=%region_23.28, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.155 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.125, %reduce.126.clone.1) } -%fused_computation.245 (param_0.681: f32[14336,4,4096]) -> bf16[4,14336,4096] { - %param_0.681 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.681), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.317 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.244 (param_0.694: f32[14336,4,4096]) -> bf16[4,14336,4096] { + %param_0.694 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) + %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.694), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} + ROOT %bitcast.323 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.246 (param_0.683: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.683 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.683), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.318 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.245 (param_0.696: f32[4096,4,14336]) -> bf16[4,4096,14336] { + %param_0.696 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.696), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} + ROOT %bitcast.324 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.247 (param_0.685: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.685 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.685), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.319 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.246 (param_0.698: f32[4096,4,14336]) -> bf16[4,4096,14336] { + %param_0.698 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.698), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} + ROOT %bitcast.325 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_52.57 (reduce_sum.289: f32[], reduce_sum.290: f32[]) -> f32[] { @@ -312,39 +312,39 @@ StackFrames ROOT %reduce_sum.219 = f32[]{:T(128)} add(%reduce_sum.217, %reduce_sum.218), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.248 (param_0.1100: f32[14336,4,4096], param_1.1263: f32[], param_2.1099: f32[], param_3.793: f32[], param_4.493: f32[14336,4,4096], param_5.415: f32[], param_6.283: f32[4,14336,4096], param_7.182: pred[], param_8.108: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { - %param_0.1100 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %param_3.793 = f32[]{:T(128)S(6)} parameter(3) - %mul.1616.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.182 = pred[]{:T(512)S(6)} parameter(7) - %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.182), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.283 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) - %bitcast.417.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.415 = f32[]{:T(128)} parameter(5) - %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.417.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.417.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.247 (param_0.1114: f32[14336,4,4096], param_1.1277: f32[], param_2.1108: f32[], param_3.797: f32[], param_4.498: f32[14336,4,4096], param_5.423: f32[], param_6.295: f32[4,14336,4096], param_7.194: pred[], param_8.112: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { + %param_0.1114 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) + %param_3.797 = f32[]{:T(128)S(6)} parameter(3) + %mul.1550.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.797), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.194 = pred[]{:T(512)S(6)} parameter(7) + %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.194), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.295 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) + %bitcast.423.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.295), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.423 = f32[]{:T(128)} parameter(5) + %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.423), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.423.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.423.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.955.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.586.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.955.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1622.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.108 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) + %mul.1556.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.112 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) %constant.959.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1623.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1621.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.108, %mul.1623.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1622.clone.1, %mul.1621.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1099 = f32[]{:T(128)S(6)} parameter(2) - %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1099), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1557.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1555.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.112, %mul.1557.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1556.clone.1, %mul.1555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1108 = f32[]{:T(128)S(6)} parameter(2) + %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1108), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.68.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %select_n.273.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.958.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1620.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1618.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1620.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.493 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) + %mul.1554.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1552.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.498 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) %constant.957.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1619.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1617.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.493, %mul.1619.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1618.clone.1, %mul.1617.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1263 = f32[]{:T(128)S(6)} parameter(1) - %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1263), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1553.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1551.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.498, %mul.1553.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1552.clone.1, %mul.1551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1277 = f32[]{:T(128)S(6)} parameter(1) + %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.783.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.819.clone.1, %div.784.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.66.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} sqrt(%div.783.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.956.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -352,13 +352,13 @@ StackFrames %add.817.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%sqrt.66.clone.1, %add.818.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.264.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%div.785.clone.1, %add.817.clone.1), metadata={op_name="multiply.34"} %div.782.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.820.clone.1, %multiply.264.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1615.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1100, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1615.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1614.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1616.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1100, %mul.1614.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.120 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%add.815.clone.1, %add.815.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1549.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1114, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1548.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1550.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1114, %mul.1548.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.171 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%add.815.clone.1, %add.815.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1001 = f32[]{:T(128)} constant(0) - %reduce.127 = f32[]{:T(128)} reduce(%square.120, %constant.1001), dimensions={0,1,2}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.127 = f32[]{:T(128)} reduce(%square.171, %constant.1001), dimensions={0,1,2}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.130.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1001), dimensions={0,1,2}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.137 = (f32[]{:T(128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.127, %add.815.clone.1, %add.819.clone.1, %add.820.clone.1, %reduce.130.clone.1) } @@ -375,39 +375,39 @@ StackFrames ROOT %reduce_sum.213 = f32[]{:T(128)} add(%reduce_sum.211, %reduce_sum.212), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.249 (param_0.1101: f32[4096,4,14336], param_1.1264: f32[], param_2.1100: f32[], param_3.794: f32[], param_4.494: f32[4096,4,14336], param_5.416: f32[], param_6.284: f32[4,4096,14336], param_7.183: pred[], param_8.109: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1101 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.794 = f32[]{:T(128)S(6)} parameter(3) - %mul.1626.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.183 = pred[]{:T(512)S(6)} parameter(7) - %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.183), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.419.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.416 = f32[]{:T(128)} parameter(5) - %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.419.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.419.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.248 (param_0.1115: f32[4096,4,14336], param_1.1278: f32[], param_2.1109: f32[], param_3.798: f32[], param_4.499: f32[4096,4,14336], param_5.424: f32[], param_6.296: f32[4,4096,14336], param_7.195: pred[], param_8.113: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { + %param_0.1115 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.798 = f32[]{:T(128)S(6)} parameter(3) + %mul.1560.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.798), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.195 = pred[]{:T(512)S(6)} parameter(7) + %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.195), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.296 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.425.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.296), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.424 = f32[]{:T(128)} parameter(5) + %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.424), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.425.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.425.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.961.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.592.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.961.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1630.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.109 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) + %mul.1564.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.113 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) %constant.965.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.591.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.965.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1629.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.109, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1630.clone.1, %mul.1629.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) - %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1563.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.113, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1564.clone.1, %mul.1563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1109 = f32[]{:T(128)S(6)} parameter(2) + %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1109), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.69.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %select_n.277.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.964.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.590.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.964.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1628.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.494 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) + %mul.1562.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.499 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) %constant.963.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.589.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.963.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1627.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.494, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1628.clone.1, %mul.1627.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1264 = f32[]{:T(128)S(6)} parameter(1) - %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1264), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1561.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.499, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1562.clone.1, %mul.1561.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1278 = f32[]{:T(128)S(6)} parameter(1) + %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1278), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.791.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.824.clone.1, %div.792.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.67.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.791.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.962.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -415,13 +415,13 @@ StackFrames %add.823.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.67.clone.1, %broadcast.587.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.265.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.793.clone.1, %add.823.clone.1), metadata={op_name="multiply.33"} %div.790.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.825.clone.1, %multiply.265.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1625.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1101, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1625.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1624.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1626.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1101, %mul.1624.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.121 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.821.clone.1, %add.821.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1559.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1115, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1558.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1560.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1115, %mul.1558.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.172 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.821.clone.1, %add.821.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1002 = f32[]{:T(128)} constant(0) - %reduce.128 = f32[]{:T(128)} reduce(%square.121, %constant.1002), dimensions={0,1,2}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.128 = f32[]{:T(128)} reduce(%square.172, %constant.1002), dimensions={0,1,2}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.131.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1002), dimensions={0,1,2}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.138 = (f32[]{:T(128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.128, %add.821.clone.1, %add.824.clone.1, %add.825.clone.1, %reduce.131.clone.1) } @@ -438,39 +438,39 @@ StackFrames ROOT %reduce_sum.210 = f32[]{:T(128)} add(%reduce_sum.205, %reduce_sum.206), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.250 (param_0.1102: f32[4096,4,14336], param_1.1265: f32[], param_2.1101: f32[], param_3.795: f32[], param_4.495: f32[4096,4,14336], param_5.417: f32[], param_6.285: f32[4,4096,14336], param_7.184: pred[], param_8.110: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1102 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.795 = f32[]{:T(128)S(6)} parameter(3) - %mul.1633.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.184 = pred[]{:T(512)S(6)} parameter(7) - %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.184), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.285 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.421.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.417 = f32[]{:T(128)} parameter(5) - %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.421.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.421.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.249 (param_0.1116: f32[4096,4,14336], param_1.1279: f32[], param_2.1110: f32[], param_3.799: f32[], param_4.500: f32[4096,4,14336], param_5.425: f32[], param_6.297: f32[4,4096,14336], param_7.196: pred[], param_8.114: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { + %param_0.1116 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.799 = f32[]{:T(128)S(6)} parameter(3) + %mul.1567.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.799), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.196 = pred[]{:T(512)S(6)} parameter(7) + %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.196), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.297 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.427.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.297), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.425 = f32[]{:T(128)} parameter(5) + %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.425), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.427.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.427.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.967.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.598.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.967.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1637.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.110 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) + %mul.1571.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.114 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) %constant.971.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.597.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.971.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1636.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.110, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1637.clone.1, %mul.1636.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) - %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1570.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.114, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1571.clone.1, %mul.1570.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1110 = f32[]{:T(128)S(6)} parameter(2) + %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1110), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.70.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %select_n.281.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.970.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.596.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.970.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1635.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.495 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) + %mul.1569.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.500 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) %constant.969.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.595.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.969.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1634.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.495, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1635.clone.1, %mul.1634.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1265 = f32[]{:T(128)S(6)} parameter(1) - %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1265), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1568.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.500, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1569.clone.1, %mul.1568.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1279 = f32[]{:T(128)S(6)} parameter(1) + %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1279), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.799.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.829.clone.1, %div.800.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.68.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.799.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.968.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -478,13 +478,13 @@ StackFrames %add.828.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.68.clone.1, %broadcast.593.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.266.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.801.clone.1, %add.828.clone.1), metadata={op_name="multiply.32"} %div.798.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.830.clone.1, %multiply.266.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1632.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1102, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1632.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1631.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1633.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1102, %mul.1631.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.122 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.826.clone.1, %add.826.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1566.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1116, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1565.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1567.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1116, %mul.1565.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.173 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.826.clone.1, %add.826.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1003 = f32[]{:T(128)} constant(0) - %reduce.129 = f32[]{:T(128)} reduce(%square.122, %constant.1003), dimensions={0,1,2}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.129 = f32[]{:T(128)} reduce(%square.173, %constant.1003), dimensions={0,1,2}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.132.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1003), dimensions={0,1,2}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.139 = (f32[]{:T(128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.129, %add.826.clone.1, %add.829.clone.1, %add.830.clone.1, %reduce.132.clone.1) } @@ -495,12 +495,12 @@ StackFrames ROOT %reduce_sum.183 = f32[]{:T(128)} add(%reduce_sum.178, %reduce_sum.182), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.268 (param_0.1106: f32[4,4096,32,128]) -> f32[] { - %param_0.1106 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.323 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1408 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.323, %bitcast.323), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.267 (param_0.1120: f32[4,4096,32,128]) -> f32[] { + %param_0.1120 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.329 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1120), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.176 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.329, %bitcast.329), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1007 = f32[]{:T(128)} constant(0) - ROOT %reduce.133 = f32[]{:T(128)} reduce(%mul.1408, %constant.1007), dimensions={0,1,2,3}, to_apply=%region_30.35, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.133 = f32[]{:T(128)} reduce(%square.176, %constant.1007), dimensions={0,1,2,3}, to_apply=%region_30.35, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_29.34 (reduce_sum.175: f32[], reduce_sum.176: f32[]) -> f32[] { @@ -509,18 +509,18 @@ StackFrames ROOT %reduce_sum.177 = f32[]{:T(128)} add(%reduce_sum.175, %reduce_sum.176), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.270 (param_0.1107: f32[4,32,128,4096]) -> f32[] { - %param_0.1107 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.327 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1411 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.327, %bitcast.327), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.269 (param_0.1121: f32[4,32,128,4096]) -> f32[] { + %param_0.1121 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.333 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1121), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.179 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.333, %bitcast.333), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1008 = f32[]{:T(128)} constant(0) - ROOT %reduce.134 = f32[]{:T(128)} reduce(%mul.1411, %constant.1008), dimensions={0,1,2,3}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.134 = f32[]{:T(128)} reduce(%square.179, %constant.1008), dimensions={0,1,2,3}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.271 (param_0.735: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { - %param_0.735 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.735), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.328 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.270 (param_0.748: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { + %param_0.748 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) + %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.748), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} + ROOT %bitcast.334 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_57.62 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { @@ -535,39 +535,39 @@ StackFrames ROOT %reduce_sum.246 = f32[]{:T(128)} add(%reduce_sum.241, %reduce_sum.245), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.272 (param_0.1095: f32[4096,4,32,128], param_1.1258: f32[], param_2.1094: f32[], param_3.788: f32[], param_4.488: f32[4096,4,32,128], param_5.410: f32[], param_6.278: f32[4,4096,32,128], param_7.177: pred[], param_8.103: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { - %param_0.1095 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.788 = f32[]{:T(128)S(6)} parameter(3) - %mul.1575.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.788), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.177 = pred[]{:T(512)S(6)} parameter(7) - %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.177), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.278 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.407.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.278), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.410 = f32[]{:T(128)} parameter(5) - %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.410), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.407.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.407.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.271 (param_0.1109: f32[4096,4,32,128], param_1.1272: f32[], param_2.1103: f32[], param_3.792: f32[], param_4.493: f32[4096,4,32,128], param_5.418: f32[], param_6.290: f32[4,4096,32,128], param_7.189: pred[], param_8.107: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { + %param_0.1109 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.792 = f32[]{:T(128)S(6)} parameter(3) + %mul.1509.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.189 = pred[]{:T(512)S(6)} parameter(7) + %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.189), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.290 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.413.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.290), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.418 = f32[]{:T(128)} parameter(5) + %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.413.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.413.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.925.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.564.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.925.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %mul.1581.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.103 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1515.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.107 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) %constant.929.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1582.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1580.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.103, %mul.1582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1581.clone.1, %mul.1580.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1094 = f32[]{:T(128)S(6)} parameter(2) - %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1094), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1516.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1514.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.107, %mul.1516.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1515.clone.1, %mul.1514.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1103 = f32[]{:T(128)S(6)} parameter(2) + %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1103), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.63.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %select_n.253.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.928.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1579.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1577.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1579.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.488 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1513.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1511.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1513.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.493 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) %constant.927.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1578.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1576.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.488, %mul.1578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1577.clone.1, %mul.1576.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1258 = f32[]{:T(128)S(6)} parameter(1) - %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1258), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1512.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1510.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.493, %mul.1512.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1511.clone.1, %mul.1510.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) + %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.743.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.792.clone.1, %div.744.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.61.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} sqrt(%div.743.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.926.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -575,13 +575,13 @@ StackFrames %add.790.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%sqrt.61.clone.1, %add.791.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.259.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%div.745.clone.1, %add.790.clone.1), metadata={op_name="multiply.39"} %div.742.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.793.clone.1, %multiply.259.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1574.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1095, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1574.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1573.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1575.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1095, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.123 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%add.788.clone.1, %add.788.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1508.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1109, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1508.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1507.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1509.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1109, %mul.1507.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.180 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%add.788.clone.1, %add.788.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.996 = f32[]{:T(128)} constant(0) - %reduce.135 = f32[]{:T(128)} reduce(%square.123, %constant.996), dimensions={0,1,2,3}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.135 = f32[]{:T(128)} reduce(%square.180, %constant.996), dimensions={0,1,2,3}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.139.clone.1 = f32[]{:T(128)} reduce(%integer_pow.63.clone.1, %constant.996), dimensions={0,1,2,3}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.140 = (f32[]{:T(128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.135, %add.788.clone.1, %add.792.clone.1, %add.793.clone.1, %reduce.139.clone.1) } @@ -598,39 +598,39 @@ StackFrames ROOT %reduce_sum.240 = f32[]{:T(128)} add(%reduce_sum.238, %reduce_sum.239), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.273 (param_0.1096: f32[32,4,128,4096], param_1.1259: f32[], param_2.1095: f32[], param_3.789: f32[], param_4.489: f32[32,4,128,4096], param_5.411: f32[], param_6.279: f32[4,32,128,4096], param_7.178: pred[], param_8.104: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { - %param_0.1096 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %param_3.789 = f32[]{:T(128)S(6)} parameter(3) - %mul.1585.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.178 = pred[]{:T(512)S(6)} parameter(7) - %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.178), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.279 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.409.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.411 = f32[]{:T(128)} parameter(5) - %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.411), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.409.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.409.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.272 (param_0.1110: f32[32,4,128,4096], param_1.1273: f32[], param_2.1104: f32[], param_3.793: f32[], param_4.494: f32[32,4,128,4096], param_5.419: f32[], param_6.291: f32[4,32,128,4096], param_7.190: pred[], param_8.108: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { + %param_0.1110 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) + %param_3.793 = f32[]{:T(128)S(6)} parameter(3) + %mul.1519.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.190 = pred[]{:T(512)S(6)} parameter(7) + %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.190), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.291 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.415.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.291), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.419 = f32[]{:T(128)} parameter(5) + %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.419), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.415.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.415.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.931.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.566.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.931.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %mul.1591.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.104 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) + %mul.1525.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.108 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) %constant.935.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1592.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1590.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.104, %mul.1592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1591.clone.1, %mul.1590.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1095 = f32[]{:T(128)S(6)} parameter(2) - %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1095), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1526.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1524.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.108, %mul.1526.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1525.clone.1, %mul.1524.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1104 = f32[]{:T(128)S(6)} parameter(2) + %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1104), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.64.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %select_n.257.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.934.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1589.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1587.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.489 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) + %mul.1523.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1521.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1523.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.494 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) %constant.933.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1588.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1586.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.489, %mul.1588.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1587.clone.1, %mul.1586.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1259 = f32[]{:T(128)S(6)} parameter(1) - %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1259), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1522.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1520.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.494, %mul.1522.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1521.clone.1, %mul.1520.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1273 = f32[]{:T(128)S(6)} parameter(1) + %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1273), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.751.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.798.clone.1, %div.752.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.62.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} sqrt(%div.751.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.932.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -638,13 +638,13 @@ StackFrames %add.796.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%sqrt.62.clone.1, %add.797.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.260.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%div.753.clone.1, %add.796.clone.1), metadata={op_name="multiply.38"} %div.750.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.799.clone.1, %multiply.260.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1584.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1096, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1584.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1583.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1585.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1096, %mul.1583.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.124 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%add.794.clone.1, %add.794.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1518.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1110, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1518.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1517.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1519.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1110, %mul.1517.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.181 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%add.794.clone.1, %add.794.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.997 = f32[]{:T(128)} constant(0) - %reduce.136 = f32[]{:T(128)} reduce(%square.124, %constant.997), dimensions={0,1,2,3}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.136 = f32[]{:T(128)} reduce(%square.181, %constant.997), dimensions={0,1,2,3}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.140.clone.1 = f32[]{:T(128)} reduce(%integer_pow.64.clone.1, %constant.997), dimensions={0,1,2,3}, to_apply=%region_42.47, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.141 = (f32[]{:T(128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.136, %add.794.clone.1, %add.798.clone.1, %add.799.clone.1, %reduce.140.clone.1) } @@ -655,23 +655,23 @@ StackFrames ROOT %reduce_sum.267 = f32[]{:T(128)} add(%reduce_sum.262, %reduce_sum.266), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.280 (param_0.1115: bf16[4,128,128256], param_1.1274: f32[4,128], param_2.1106: s32[4,128], param_3.799: bf16[4,128]) -> f32[4,128] { - %param_2.1106 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1106), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.279 (param_0.1129: bf16[4,128,128256], param_1.1288: f32[4,128], param_2.1115: s32[4,128], param_3.803: bf16[4,128]) -> f32[4,128] { + %param_2.1115 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1115), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.25 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.24 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1115 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.938 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1115), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.799 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.799), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.938, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1274 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_0.1129 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.950 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1129), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_3.803 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) + %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.803), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.950, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.60 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %constant.1017 = f32[]{:T(128)} constant(0) %broadcast.511 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%constant.1017), dimensions={}, metadata={op_name="broadcast.83"} - %mul.1424 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1424, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1373 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1373, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_7.10 (reduce_sum.93: f32[], reduce_sum.94: f32[]) -> f32[] { @@ -680,12 +680,12 @@ StackFrames ROOT %reduce_sum.95 = f32[]{:T(128)} add(%reduce_sum.93, %reduce_sum.94), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.285 (param_0.1116: bf16[4,128,128256], param_1.1275: bf16[4,128]) -> f32[4,128] { - %param_0.1116 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.944 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1116), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1275 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1275), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.944, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.284 (param_0.1130: bf16[4,128,128256], param_1.1289: bf16[4,128]) -> f32[4,128] { + %param_0.1130 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.956 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1130), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1289 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1289), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.956, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.54 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} %constant.1018 = f32[]{:T(128)} constant(0) ROOT %reduce.138 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1018), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} @@ -703,23 +703,23 @@ StackFrames ROOT %reduce_sum.171 = f32[]{:T(128)} add(%reduce_sum.169, %reduce_sum.170), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.291 (param_0.1108: f32[4,4096,8,128], param_1.1268: f32[4,4096,8,128]) -> (f32[], f32[]) { - %param_0.1108 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.344 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1439 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.344, %bitcast.344), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.290 (param_0.1122: f32[4,4096,8,128], param_1.1282: f32[4,4096,8,128]) -> (f32[], f32[]) { + %param_0.1122 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.350 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1122), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.184 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.350, %bitcast.350), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1009 = f32[]{:T(128)} constant(0) - %reduce.141 = f32[]{:T(128)} reduce(%mul.1439, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1268 = f32[4,4096,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) - %bitcast.348.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1442.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.348.clone.1, %bitcast.348.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.142.clone.1 = f32[]{:T(128)} reduce(%mul.1442.clone.1, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.141 = f32[]{:T(128)} reduce(%square.184, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1282 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(1) + %bitcast.354.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.187.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.354.clone.1, %bitcast.354.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.142.clone.1 = f32[]{:T(128)} reduce(%square.187.clone.1, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.156 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.141, %reduce.142.clone.1) } -%fused_computation.294 (param_0.794: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.794 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.794), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.349 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.293 (param_0.807: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { + %param_0.807 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.807), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} + ROOT %bitcast.355 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_58.63 (reduce_sum.324: f32[], reduce_sum.325: f32[]) -> f32[] { @@ -734,39 +734,39 @@ StackFrames ROOT %reduce_sum.252 = f32[]{:T(128)} add(%reduce_sum.247, %reduce_sum.248), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.295 (param_0.1094: f32[4096,4,8,128], param_1.1257: f32[], param_2.1093: f32[], param_3.787: f32[], param_4.487: f32[4096,4,8,128], param_5.409: f32[], param_6.277: f32[4,4096,8,128], param_7.176: pred[], param_8.102: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1094 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %param_3.787 = f32[]{:T(128)S(6)} parameter(3) - %mul.1568.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.787), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.176 = pred[]{:T(512)S(6)} parameter(7) - %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.176), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.277 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.405.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.277), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.409 = f32[]{:T(128)} parameter(5) - %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.405.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.405.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.294 (param_0.1108: f32[4096,4,8,128], param_1.1271: f32[], param_2.1102: f32[], param_3.791: f32[], param_4.492: f32[4096,4,8,128], param_5.417: f32[], param_6.289: f32[4,4096,8,128], param_7.188: pred[], param_8.106: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { + %param_0.1108 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.791 = f32[]{:T(128)S(6)} parameter(3) + %mul.1502.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.188 = pred[]{:T(512)S(6)} parameter(7) + %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.188), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.289 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.289), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.417 = f32[]{:T(128)} parameter(5) + %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.411.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.919.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.562.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.919.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1572.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1506.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.106 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.923.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.561.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.923.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1571.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.102, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1572.clone.1, %mul.1571.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1093 = f32[]{:T(128)S(6)} parameter(2) - %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1093), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1505.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.106, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1506.clone.1, %mul.1505.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) + %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.62.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %select_n.249.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.922.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.560.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.922.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1570.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.487 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1504.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.492 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.921.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.559.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.921.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1569.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.487, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1570.clone.1, %mul.1569.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1257 = f32[]{:T(128)S(6)} parameter(1) - %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1257), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1503.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.492, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1504.clone.1, %mul.1503.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1271 = f32[]{:T(128)S(6)} parameter(1) + %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1271), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.735.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.786.clone.1, %div.736.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.60.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.735.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.920.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -774,15 +774,15 @@ StackFrames %add.785.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.60.clone.1, %broadcast.557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.258.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.737.clone.1, %add.785.clone.1), metadata={op_name="multiply.40"} %div.734.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.787.clone.1, %multiply.258.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1567.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1094, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1566.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1568.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1094, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.125 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.783.clone.1, %add.783.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1501.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1108, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1501.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1500.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1502.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1108, %mul.1500.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.188 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.783.clone.1, %add.783.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.995 = f32[]{:T(128)} constant(0) - %reduce.143 = f32[]{:T(128)} reduce(%square.125, %constant.995), dimensions={0,1,2,3}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.143 = f32[]{:T(128)} reduce(%square.188, %constant.995), dimensions={0,1,2,3}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.145.clone.1 = f32[]{:T(128)} reduce(%integer_pow.62.clone.1, %constant.995), dimensions={0,1,2,3}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) + ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) } %region_55.60 (reduce_sum.304: f32[], reduce_sum.308: f32[]) -> f32[] { @@ -797,39 +797,39 @@ StackFrames ROOT %reduce_sum.234 = f32[]{:T(128)} add(%reduce_sum.232, %reduce_sum.233), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.296 (param_0.1097: f32[4096,4,8,128], param_1.1260: f32[], param_2.1096: f32[], param_3.790: f32[], param_4.490: f32[4096,4,8,128], param_5.412: f32[], param_6.280: f32[4,4096,8,128], param_7.179: pred[], param_8.105: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1097 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.790 = f32[]{:T(128)S(6)} parameter(3) - %mul.1595.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.179 = pred[]{:T(512)S(6)} parameter(7) - %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.179), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.280 = f32[4,4096,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.412 = f32[]{:T(128)} parameter(5) - %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.411.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.295 (param_0.1111: f32[4096,4,8,128], param_1.1274: f32[], param_2.1105: f32[], param_3.794: f32[], param_4.495: f32[4096,4,8,128], param_5.420: f32[], param_6.292: f32[4,4096,8,128], param_7.191: pred[], param_8.109: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { + %param_0.1111 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.794 = f32[]{:T(128)S(6)} parameter(3) + %mul.1529.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.191 = pred[]{:T(512)S(6)} parameter(7) + %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.191), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.292 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.417.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.292), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.420 = f32[]{:T(128)} parameter(5) + %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.420), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.417.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.417.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.937.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.572.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.937.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1599.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.105 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1533.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.109 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.941.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.571.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.941.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1598.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.105, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1599.clone.1, %mul.1598.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1096 = f32[]{:T(128)S(6)} parameter(2) - %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1096), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1532.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.109, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1533.clone.1, %mul.1532.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1105 = f32[]{:T(128)S(6)} parameter(2) + %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1105), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.65.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %select_n.261.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.940.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.570.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.940.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1597.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.490 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1531.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.495 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.939.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.569.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.939.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1596.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.490, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1597.clone.1, %mul.1596.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1260 = f32[]{:T(128)S(6)} parameter(1) - %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1260), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1530.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.495, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1531.clone.1, %mul.1530.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1274 = f32[]{:T(128)S(6)} parameter(1) + %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.759.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.803.clone.1, %div.760.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.63.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.759.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.938.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -837,33 +837,33 @@ StackFrames %add.802.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.761.clone.1, %add.802.clone.1), metadata={op_name="multiply.37"} %div.758.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.804.clone.1, %multiply.261.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1594.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1097, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1594.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1593.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1595.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1097, %mul.1593.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.126 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.800.clone.1, %add.800.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1528.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1111, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1528.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1527.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1529.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1111, %mul.1527.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.189 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.800.clone.1, %add.800.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.998 = f32[]{:T(128)} constant(0) - %reduce.144 = f32[]{:T(128)} reduce(%square.126, %constant.998), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.144 = f32[]{:T(128)} reduce(%square.189, %constant.998), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.146.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.998), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.143 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.144, %add.800.clone.1, %add.803.clone.1, %add.804.clone.1, %reduce.146.clone.1) } -%fused_computation.312 (param_0.859: bf16[4,128,4096], param_1.928: f32[4,128], param_2.717: f32[4,128], param_3.448: bf16[4,128,4096], param_4.266: bf16[4096]) -> bf16[4,128,4096] { - %param_3.448 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.266 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.371 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.266), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.361 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.448, %dot_general.371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.961 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.717 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1480 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.717), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1472 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.961, %mul.1480), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.859 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.972 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.859), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.928 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1479 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.928), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1478 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.972, %mul.1479), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1472, %mul.1478), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} - ROOT %convert_element_type.959 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.311 (param_0.872: bf16[4,128,4096], param_1.941: f32[4,128], param_2.726: f32[4,128], param_3.452: bf16[4,128,4096], param_4.271: bf16[4096]) -> bf16[4,128,4096] { + %param_3.452 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.271 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %dot_general.375 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.271), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.365 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.452, %dot_general.375), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.973 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_2.726 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %mul.1423 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.726), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1415 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.973, %mul.1423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_0.872 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.984 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.872), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.941 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1422 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.941), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1421 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.984, %mul.1422), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1415, %mul.1421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} + ROOT %convert_element_type.971 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} } %region_5.8 (reduce_sum.87: f32[], reduce_sum.88: f32[]) -> f32[] { @@ -872,12 +872,12 @@ StackFrames ROOT %reduce_sum.92 = f32[]{:T(128)} add(%reduce_sum.87, %reduce_sum.88), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.313 (param_0.1117: bf16[4,128,4096]) -> f32[4,128] { - %param_0.1117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.963 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1117), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.129 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.963, %convert_element_type.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} +%fused_computation.312 (param_0.1131: bf16[4,128,4096]) -> f32[4,128] { + %param_0.1131 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.975 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1131), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.192 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.975, %convert_element_type.975), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} %constant.1019 = f32[]{:T(128)} constant(0) - ROOT %reduce.147 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.129, %constant.1019), dimensions={2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + ROOT %reduce.147 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.192, %constant.1019), dimensions={2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_10.13 (reduce_sum.102: f32[], reduce_sum.106: f32[]) -> f32[] { @@ -886,17 +886,17 @@ StackFrames ROOT %reduce_sum.107 = f32[]{:T(128)} add(%reduce_sum.102, %reduce_sum.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.315 (param_0.1112: bf16[4,128,4096], param_1.1271: bf16[4,128,4096], param_2.1104: bf16[4096]) -> f32[4,128] { - %param_0.1112 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.970 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1112), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1104 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.370 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1104), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.360 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1271, %dot_general.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.969 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1476 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.970, %convert_element_type.969), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} +%fused_computation.314 (param_0.1126: bf16[4,128,4096], param_1.1285: bf16[4,128,4096], param_2.1113: bf16[4096]) -> f32[4,128] { + %param_0.1126 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1126), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1285 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.374 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1113), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.364 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1285, %dot_general.374), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.981 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.1419 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %convert_element_type.981), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.1013 = f32[]{:T(128)} constant(0) - ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1476, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1419, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_8.11 (dot_general.182: bf16[], dot_general.183: bf16[]) -> bf16[] { @@ -905,86 +905,86 @@ StackFrames ROOT %add.168 = bf16[]{:T(256)} add(%dot_general.182, %dot_general.183), metadata={op_name="add.54"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.236.clone.clone (param_0.1081: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1081 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1021 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1081), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.443 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1021), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +%fused_computation.235.clone.clone (param_0.1095: f32[4096,128256]) -> bf16[4096,128256,1] { + %param_0.1095 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %convert_element_type.1033 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1095), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + ROOT %bitcast.449 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} } -%fused_computation.281.clone.1.clone.clone (param_0.1082: bf16[4,128,128256], param_1.1247: s32[4,128], param_2.1072: f32[4,128], param_3.778: f32[4,128], param_4.479: bf16[4,128], param_5.401: f32[4,128]) -> bf16[4,128,128256] { - %param_5.401 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1669 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.401), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.778 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1668 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.778), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1082 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1024 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.479 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.479), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1024, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.280.clone.1.clone.clone (param_0.1096: bf16[4,128,128256], param_1.1261: s32[4,128], param_2.1081: f32[4,128], param_3.782: f32[4,128], param_4.484: bf16[4,128], param_5.409: f32[4,128]) -> bf16[4,128,128256] { + %param_5.409 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.1603 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.782 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.1602 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.782), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1096 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1036 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1096), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.484 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.484), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1036, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.60 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1667 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1668, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1072 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1072), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1667, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1247 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1247), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.1601 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1602, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1081), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1601, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1261 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1261), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.42 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.41 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1023 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1023), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1666 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1669, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1022 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1666), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.316 (param_0.1080: f32[4,128], param_1.1246: bf16[4,128,4096], param_2.1073: f32[4096,128256], param_3.779: bf16[4,128,128256], param_4.480: s32[4,128], param_5.402: f32[4,128], param_6.272: f32[4,128], param_7.171: bf16[4,128], param_8.98: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { - %param_3.779 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.480 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.272 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.171 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.779, %param_4.480, %param_5.402, %param_6.272, %param_7.171, /*index=5*/%param_8.98), kind=kLoop, calls=%fused_computation.281.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1073 = f32[4096,128256]{1,0:T(8,128)} parameter(2) - %fusion.209.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1073), kind=kLoop, calls=%fused_computation.236.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.80.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.209.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %param_1.1246 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1246), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1080 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1491 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1080), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1490 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %mul.1491), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.981 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1490), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.80.clone.1, %convert_element_type.981), metadata={op_name="multiply.206"} + %convert_element_type.1035 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1035), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1600 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1603, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1034 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1600), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.315 (param_0.1094: f32[4,128], param_1.1260: bf16[4,128,4096], param_2.1082: f32[4096,128256], param_3.783: bf16[4,128,128256], param_4.485: s32[4,128], param_5.410: f32[4,128], param_6.284: f32[4,128], param_7.183: bf16[4,128], param_8.102: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { + %param_3.783 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) + %param_4.485 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %param_5.410 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.284 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.183 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) + %param_8.102 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) + %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.783, %param_4.485, %param_5.410, %param_6.284, %param_7.183, /*index=5*/%param_8.102), kind=kLoop, calls=%fused_computation.280.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_2.1082 = f32[4096,128256]{1,0:T(8,128)} parameter(2) + %fusion.219.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1082), kind=kLoop, calls=%fused_computation.235.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %convolution.86.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.219.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %param_1.1260 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.994 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1260), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1094 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1434 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1094), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1433 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.994, %mul.1434), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.993 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1433), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.86.clone.1, %convert_element_type.993), metadata={op_name="multiply.206"} %constant.874 = bf16[]{:T(256)} constant(0) %reduce.149 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.252, %constant.874), dimensions={0,1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.80.clone.1) + ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.86.clone.1) } -%fused_computation.324 (param_0.891: f32[64], param_1.961: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.961 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.961), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.891 = f32[64]{0:T(128)S(1)} parameter(0) - %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.891), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} +%fused_computation.323 (param_0.904: f32[64], param_1.974: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.974), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %param_0.904 = f32[64]{0:T(128)S(1)} parameter(0) + %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.904), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %div.618 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.621, %div.619), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.618), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} - %convert_element_type.990 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1002 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %cos.41.clone.1 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.618), metadata={op_name="jit(train_step)/layers/cos" stack_frame_id=0} - %convert_element_type.989.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.990, %convert_element_type.989.clone.1) + %convert_element_type.1001.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1002, %convert_element_type.1001.clone.1) } -%fused_computation.325 (param_0.888: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.888 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.324 (param_0.901: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.901 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.866 = bf16[]{:T(256)} constant(-inf) - %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.888, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.888, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.34 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.38, %pad.37), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } -%fused_computation.326 (param_0.890: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.890 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.325 (param_0.903: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.903 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.865 = bf16[]{:T(256)} constant(-inf) - %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.890, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.890, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.35 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.40, %pad.39), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } @@ -1000,16 +1000,16 @@ StackFrames ROOT %reduce_sum.162 = f32[]{:T(128)} add(%reduce_sum.157, %reduce_sum.161), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.330 (param_0.1109: f32[4,4096], param_1.1269: f32[4,4096]) -> (f32[], f32[]) { - %param_0.1109 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.365 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1500 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.365, %bitcast.365), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.329 (param_0.1123: f32[4,4096], param_1.1283: f32[4,4096]) -> (f32[], f32[]) { + %param_0.1123 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.371 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.195 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1010 = f32[]{:T(128)} constant(0) - %reduce.150 = f32[]{:T(128)} reduce(%mul.1500, %constant.1010), dimensions={0,1}, to_apply=%region_27.32, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1269 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) - %bitcast.369.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1503.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.369.clone.1, %bitcast.369.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.151.clone.1 = f32[]{:T(128)} reduce(%mul.1503.clone.1, %constant.1010), dimensions={0,1}, to_apply=%region_26.31, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.150 = f32[]{:T(128)} reduce(%square.195, %constant.1010), dimensions={0,1}, to_apply=%region_27.32, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1283 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) + %bitcast.375.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.198.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.151.clone.1 = f32[]{:T(128)} reduce(%square.198.clone.1, %constant.1010), dimensions={0,1}, to_apply=%region_26.31, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.150, %reduce.151.clone.1) } @@ -1025,39 +1025,39 @@ StackFrames ROOT %reduce_sum.231 = f32[]{:T(128)} add(%reduce_sum.226, %reduce_sum.227), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.333 (param_0.1098: f32[4096,4], param_1.1261: f32[], param_2.1097: f32[], param_3.791: f32[], param_4.491: f32[4096,4], param_5.413: f32[], param_6.281: f32[4,4096], param_7.180: pred[], param_8.106: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1098 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.791 = f32[]{:T(128)S(6)} parameter(3) - %mul.1602.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.180 = pred[]{:T(512)S(6)} parameter(7) - %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.180), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.281 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.413.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.281), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.413 = f32[]{:T(128)} parameter(5) - %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.413), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.413.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.413.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.332 (param_0.1112: f32[4096,4], param_1.1275: f32[], param_2.1106: f32[], param_3.795: f32[], param_4.496: f32[4096,4], param_5.421: f32[], param_6.293: f32[4,4096], param_7.192: pred[], param_8.110: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1112 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.795 = f32[]{:T(128)S(6)} parameter(3) + %mul.1536.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.192 = pred[]{:T(512)S(6)} parameter(7) + %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.192), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.293 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.419.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.293), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.421 = f32[]{:T(128)} parameter(5) + %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.421), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.419.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.419.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.943.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.578.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.943.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1606.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.106 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1540.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.110 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) %constant.947.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.577.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.947.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1605.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.106, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1606.clone.1, %mul.1605.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1097 = f32[]{:T(128)S(6)} parameter(2) - %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1097), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1539.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.110, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1540.clone.1, %mul.1539.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1106 = f32[]{:T(128)S(6)} parameter(2) + %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1106), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.66.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %select_n.265.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.946.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.576.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.946.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1604.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.491 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1538.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.496 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) %constant.945.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.575.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.945.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1603.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.491, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1604.clone.1, %mul.1603.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1261 = f32[]{:T(128)S(6)} parameter(1) - %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1261), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1537.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.496, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1538.clone.1, %mul.1537.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1275 = f32[]{:T(128)S(6)} parameter(1) + %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1275), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.767.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.808.clone.1, %div.768.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.64.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.767.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.944.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1065,13 +1065,13 @@ StackFrames %add.807.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.262.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.769.clone.1, %add.807.clone.1), metadata={op_name="multiply.36"} %div.766.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.809.clone.1, %multiply.262.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1601.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1098, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1601.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1600.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1602.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1098, %mul.1600.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.130 = f32[4096,4]{0,1:T(4,128)} multiply(%add.805.clone.1, %add.805.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1535.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1112, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1535.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1534.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1536.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1112, %mul.1534.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.199 = f32[4096,4]{0,1:T(4,128)} multiply(%add.805.clone.1, %add.805.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.999 = f32[]{:T(128)} constant(0) - %reduce.152 = f32[]{:T(128)} reduce(%square.130, %constant.999), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.152 = f32[]{:T(128)} reduce(%square.199, %constant.999), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.154.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.999), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.144 = (f32[]{:T(128)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.152, %add.805.clone.1, %add.808.clone.1, %add.809.clone.1, %reduce.154.clone.1) } @@ -1088,39 +1088,39 @@ StackFrames ROOT %reduce_sum.225 = f32[]{:T(128)} add(%reduce_sum.220, %reduce_sum.224), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.334 (param_0.1099: f32[4096,4], param_1.1262: f32[], param_2.1098: f32[], param_3.792: f32[], param_4.492: f32[4096,4], param_5.414: f32[], param_6.282: f32[4,4096], param_7.181: pred[], param_8.107: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1099 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.792 = f32[]{:T(128)S(6)} parameter(3) - %mul.1609.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.181 = pred[]{:T(512)S(6)} parameter(7) - %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.181), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.282 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.415.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.414 = f32[]{:T(128)} parameter(5) - %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.414), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.415.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.415.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.333 (param_0.1113: f32[4096,4], param_1.1276: f32[], param_2.1107: f32[], param_3.796: f32[], param_4.497: f32[4096,4], param_5.422: f32[], param_6.294: f32[4,4096], param_7.193: pred[], param_8.111: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1113 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.796 = f32[]{:T(128)S(6)} parameter(3) + %mul.1543.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.193 = pred[]{:T(512)S(6)} parameter(7) + %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.193), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.294 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.421.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.294), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.422 = f32[]{:T(128)} parameter(5) + %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.422), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.421.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.421.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.949.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.584.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.949.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1613.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.107 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1547.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.111 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) %constant.953.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.583.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.953.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1612.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.107, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1613.clone.1, %mul.1612.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1098 = f32[]{:T(128)S(6)} parameter(2) - %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1098), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1546.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.111, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1547.clone.1, %mul.1546.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1107 = f32[]{:T(128)S(6)} parameter(2) + %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1107), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.67.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %select_n.269.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.952.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.582.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.952.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1611.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.492 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1545.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.497 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) %constant.951.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.581.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.951.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1610.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.492, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1611.clone.1, %mul.1610.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1262 = f32[]{:T(128)S(6)} parameter(1) - %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1262), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1544.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.497, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1545.clone.1, %mul.1544.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1276 = f32[]{:T(128)S(6)} parameter(1) + %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1276), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.775.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.813.clone.1, %div.776.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.65.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.775.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.950.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1128,13 +1128,13 @@ StackFrames %add.812.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.65.clone.1, %broadcast.579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.263.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.777.clone.1, %add.812.clone.1), metadata={op_name="multiply.35"} %div.774.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.814.clone.1, %multiply.263.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1608.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1099, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1608.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1607.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1609.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1099, %mul.1607.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.131 = f32[4096,4]{0,1:T(4,128)} multiply(%add.810.clone.1, %add.810.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1542.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1113, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1542.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1541.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1543.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1113, %mul.1541.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.200 = f32[4096,4]{0,1:T(4,128)} multiply(%add.810.clone.1, %add.810.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1000 = f32[]{:T(128)} constant(0) - %reduce.153 = f32[]{:T(128)} reduce(%square.131, %constant.1000), dimensions={0,1}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.153 = f32[]{:T(128)} reduce(%square.200, %constant.1000), dimensions={0,1}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.155.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1000), dimensions={0,1}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.145 = (f32[]{:T(128)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.153, %add.810.clone.1, %add.813.clone.1, %add.814.clone.1, %reduce.155.clone.1) } @@ -1145,12 +1145,12 @@ StackFrames ROOT %reduce_sum.101 = f32[]{:T(128)} add(%reduce_sum.99, %reduce_sum.100), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.345 (param_0.1113: bf16[4096]) -> f32[] { - %param_0.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.994 = f32[4096]{0:T(1024)} convert(%param_0.1113), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1520 = f32[4096]{0:T(1024)} multiply(%convert_element_type.994, %convert_element_type.994), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.344 (param_0.1127: bf16[4096]) -> f32[] { + %param_0.1127 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.1006 = f32[4096]{0:T(1024)} convert(%param_0.1127), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.203 = f32[4096]{0:T(1024)} multiply(%convert_element_type.1006, %convert_element_type.1006), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1014 = f32[]{:T(128)} constant(0) - ROOT %reduce.156 = f32[]{:T(128)} reduce(%mul.1520, %constant.1014), dimensions={0}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.156 = f32[]{:T(128)} reduce(%square.203, %constant.1014), dimensions={0}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_49.54 (reduce_sum.274: f32[], reduce_sum.275: f32[]) -> f32[] { @@ -1165,39 +1165,39 @@ StackFrames ROOT %reduce_sum.204 = f32[]{:T(128)} add(%reduce_sum.199, %reduce_sum.203), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.346 (param_0.1103: f32[4096], param_1.1266: f32[], param_2.1102: f32[], param_3.796: f32[], param_4.496: f32[4096], param_5.418: f32[], param_6.286: bf16[4096], param_7.185: pred[], param_8.111: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { - %param_0.1103 = f32[4096]{0:T(1024)S(1)} parameter(0) - %param_3.796 = f32[]{:T(128)S(6)} parameter(3) - %mul.1640.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.185 = pred[]{:T(512)S(6)} parameter(7) - %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.185), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.286 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1009.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.286), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.418 = f32[]{:T(128)} parameter(5) - %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1009.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1009.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.345 (param_0.1117: f32[4096], param_1.1280: f32[], param_2.1111: f32[], param_3.800: f32[], param_4.501: f32[4096], param_5.426: f32[], param_6.298: bf16[4096], param_7.197: pred[], param_8.115: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { + %param_0.1117 = f32[4096]{0:T(1024)S(1)} parameter(0) + %param_3.800 = f32[]{:T(128)S(6)} parameter(3) + %mul.1574.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.800), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.197 = pred[]{:T(512)S(6)} parameter(7) + %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.298 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1021.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.298), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_5.426 = f32[]{:T(128)} parameter(5) + %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.426), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1021.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1021.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.973.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.600.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.973.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.1646.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.111 = f32[4096]{0:T(1024)S(1)} parameter(8) + %mul.1580.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.115 = f32[4096]{0:T(1024)S(1)} parameter(8) %constant.977.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1647.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1645.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.111, %mul.1647.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1646.clone.1, %mul.1645.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) - %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1581.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1579.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.115, %mul.1581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1580.clone.1, %mul.1579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1111 = f32[]{:T(128)S(6)} parameter(2) + %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1111), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.71.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %select_n.285.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.976.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1644.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1642.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1644.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.496 = f32[4096]{0:T(1024)S(1)} parameter(4) + %mul.1578.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1576.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.501 = f32[4096]{0:T(1024)S(1)} parameter(4) %constant.975.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1643.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1641.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.496, %mul.1643.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1642.clone.1, %mul.1641.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1266 = f32[]{:T(128)S(6)} parameter(1) - %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1266), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1577.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1575.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.501, %mul.1577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1576.clone.1, %mul.1575.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1280 = f32[]{:T(128)S(6)} parameter(1) + %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1280), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.807.clone.1 = f32[4096]{0:T(1024)} divide(%add.835.clone.1, %div.808.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.69.clone.1 = f32[4096]{0:T(1024)} sqrt(%div.807.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.974.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1205,38 +1205,38 @@ StackFrames %add.833.clone.1 = f32[4096]{0:T(1024)} add(%sqrt.69.clone.1, %add.834.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.267.clone.1 = f32[4096]{0:T(1024)} multiply(%div.809.clone.1, %add.833.clone.1), metadata={op_name="multiply.31"} %div.806.clone.1 = f32[4096]{0:T(1024)} divide(%add.836.clone.1, %multiply.267.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1639.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1103, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1639.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1638.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1640.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1103, %mul.1638.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.132 = f32[4096]{0:T(1024)} multiply(%add.831.clone.1, %add.831.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.1573.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1117, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1572.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1574.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1117, %mul.1572.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.204 = f32[4096]{0:T(1024)} multiply(%add.831.clone.1, %add.831.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1004 = f32[]{:T(128)} constant(0) - %reduce.157 = f32[]{:T(128)} reduce(%square.132, %constant.1004), dimensions={0}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.157 = f32[]{:T(128)} reduce(%square.204, %constant.1004), dimensions={0}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.158.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1004), dimensions={0}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.148 = (f32[]{:T(128)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.157, %add.831.clone.1, %add.835.clone.1, %add.836.clone.1, %reduce.158.clone.1) } -%fused_computation.352 (param_0.951: s32[512]) -> s32[1024] { +%fused_computation.351 (param_0.964: s32[512]) -> s32[1024] { %constant.801 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.539 = s32[1024]{0:T(1024)} broadcast(%constant.801), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.951 = s32[512]{0:T(512)S(1)} parameter(0) + %param_0.964 = s32[512]{0:T(512)S(1)} parameter(0) %constant.802 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.951, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.964, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %constant.800 = s32[] constant(128255), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.538 = s32[1024]{0:T(1024)} broadcast(%constant.800), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.539, %pad.41, %broadcast.538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } -%fused_computation.353 (param_0.950: s32[4,128]) -> s32[512] { - %param_0.950 = s32[4,128]{1,0:T(4,128)} parameter(0) +%fused_computation.352 (param_0.963: s32[4,128]) -> s32[512] { + %param_0.963 = s32[4,128]{1,0:T(4,128)} parameter(0) %constant.888 = s32[]{:T(128)} constant(0) %broadcast.546 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.888), dimensions={}, metadata={op_name="broadcast.81"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.950, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} + %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.963, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} %constant.875 = s32[]{:T(128)} constant(128256) %add.760 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.875), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.950, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} - ROOT %bitcast.370 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.963, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} + %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} + ROOT %bitcast.376 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} } %region_61.66 (reduce_sum.345: f32[], reduce_sum.346: f32[]) -> f32[] { @@ -1251,52 +1251,52 @@ StackFrames ROOT %reduce_sum.273 = f32[]{:T(128)} add(%reduce_sum.268, %reduce_sum.269), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.354 (param_0.1114: bf16[4,128], param_1.1273: f32[4,128], param_2.1105: f32[4,128], param_3.798: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { - %param_3.798 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) +%fused_computation.353 (param_0.1128: bf16[4,128], param_1.1287: f32[4,128], param_2.1114: f32[4,128], param_3.802: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { + %param_3.802 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) %constant.979.clone.1 = s32[]{:T(128)} constant(0) %broadcast.601.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.979.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.798, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1273), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1114 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1114), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.802, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %param_1.1287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1287), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_0.1128 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1128), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} %add.762 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %square.135 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %add.762), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %square.207 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %add.762), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} %constant.1016 = f32[]{:T(128)} constant(0) %broadcast.543 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1016), dimensions={}, metadata={op_name="broadcast.32"} - %mul.1539 = f32[4,128]{1,0:T(4,128)} multiply(%square.135, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %mul.1531 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1539, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.159 = f32[]{:T(128)} reduce(%mul.1531, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1105 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1105), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} - %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1539), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %mul.1532.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1532.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %mul.1537.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1473 = f32[4,128]{1,0:T(4,128)} multiply(%square.207, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1465 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1473, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.159 = f32[]{:T(128)} reduce(%mul.1465, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_2.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1114), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1473), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1466.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1466.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1471.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %constant.891.clone.1 = f32[]{:T(128)} constant(1) %add.757.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.891.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1537.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1471.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} ROOT %tuple.149 = (f32[]{:T(128)}, f32[]{:T(128)}, pred[4,128]{1,0:T(4,128)(4,1)S(1)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%reduce.159, %reduce.160.clone.1, %ne.6.clone.1, %add.750.clone.1) } -%fused_computation.357 (param_0.974: f32[4,128], param_1.1088: f32[4,128]) -> f32[4,128] { - %param_0.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1088 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.356 (param_0.987: f32[4,128], param_1.1101: f32[4,128]) -> f32[4,128] { + %param_0.987 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1101 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.869 = f32[]{:T(128)} constant(0.000244140625) %broadcast.549 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.869), dimensions={}, metadata={op_name="broadcast.264"} - %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1088, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1101, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.867 = f32[]{:T(128)} constant(1e-05) %add.770 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.867), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.769 = f32[4,128]{1,0:T(4,128)} add(%div.656, %add.770), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %rsqrt.90 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} %div.649 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.90, %add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.864 = f32[]{:T(128)} constant(-0.5) - %mul.1543 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1536 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1543), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1535 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.974, %mul.1536), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1477 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1470 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1477), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1469 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.987, %mul.1470), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.863 = f32[]{:T(128)} constant(0.00048828125) - %mul.1542 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - ROOT %mul.1534 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1535, %mul.1542), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1476 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + ROOT %mul.1468 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1469, %mul.1476), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} } %region_0.1 (reduce_sum.67: s32[], reduce_sum.71: s32[]) -> s32[] { @@ -1305,64 +1305,64 @@ StackFrames ROOT %reduce_sum.72 = s32[]{:T(128)} add(%reduce_sum.67, %reduce_sum.71), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.361 (param_0.991: pred[4,128]) -> s32[] { - %param_0.991 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1001 = s32[4,128]{1,0:T(4,128)} convert(%param_0.991), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} +%fused_computation.360 (param_0.1004: pred[4,128]) -> s32[] { + %param_0.1004 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %convert_element_type.1013 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1004), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} %constant.889 = s32[]{:T(128)} constant(0) - ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1001, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1013, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.362 (param_0.976: f32[4,128]) -> f32[4,128] { - %param_0.976 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.361 (param_0.989: f32[4,128]) -> f32[4,128] { + %param_0.989 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.870 = f32[]{:T(128)} constant(0.000244140625) %broadcast.541 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.870), dimensions={}, metadata={op_name="broadcast.264"} - %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.976, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.989, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.868 = f32[]{:T(128)} constant(1e-05) %add.759 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.868), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.756 = f32[4,128]{1,0:T(4,128)} add(%div.654, %add.759), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} ROOT %rsqrt.88 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.756), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} } -%fused_computation.363 (param_0.977: pred[4,128], param_1.1272: f32[]) -> f32[4,128] { - %param_0.977 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} +%fused_computation.362 (param_0.990: pred[4,128], param_1.1286: f32[]) -> f32[4,128] { + %param_0.990 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %param_1.1286 = f32[]{:T(128)S(6)} parameter(1) + %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1286), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} %constant.1015 = f32[]{:T(128)} constant(0) %broadcast.545 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1015), dimensions={}, metadata={op_name="broadcast.32"} - ROOT %mul.1544 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.977, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %mul.1478 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.990, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} } -%fused_computation.365 () -> f32[64] { +%fused_computation.364 () -> f32[64] { %constant.873 = f32[]{:T(128)} constant(500000) %broadcast.552 = f32[64]{0:T(128)} broadcast(%constant.873), dimensions={}, metadata={op_name="broadcast.255"} %iota.46 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/layers/iota" stack_frame_id=0} %constant.872 = s32[]{:T(128)} constant(2) %broadcast.551 = s32[64]{0:T(128)} broadcast(%constant.872), dimensions={}, metadata={op_name="broadcast.256"} - %mul.1545 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} - %convert_element_type.1002 = f32[64]{0:T(128)} convert(%mul.1545), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %mul.1479 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} + %convert_element_type.1014 = f32[64]{0:T(128)} convert(%mul.1479), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %constant.871 = f32[]{:T(128)} constant(0.0078125) %broadcast.550 = f32[64]{0:T(128)} broadcast(%constant.871), dimensions={}, metadata={op_name="broadcast.257"} - %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1002, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1014, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.552, %div.657), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} } -%fused_computation.366 (param_0.989: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.989 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1003 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.989), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %bitcast.371 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.371, %convert_element_type.1003) +%fused_computation.365 (param_0.1002: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { + %param_0.1002 = s32[4,128]{1,0:T(4,128)} parameter(0) + %convert_element_type.1015 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1002), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %bitcast.377 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1015), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.377, %convert_element_type.1015) } -%fused_computation.369 (param_0.1089: f32[4096,4]) -> bf16[4,4096] { - %param_0.1089 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.445 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1089), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.445) +%fused_computation.369 (param_0.1103: f32[4096,4]) -> bf16[4,4096] { + %param_0.1103 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.451 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1103), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.451) } -%fused_computation.370 (param_0.1090: f32[4096,4]) -> bf16[4,4096] { - %param_0.1090 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.446 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.446) +%fused_computation.370 (param_0.1104: f32[4096,4]) -> bf16[4,4096] { + %param_0.1104 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.452 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1104), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.452) } %region_6.9 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { @@ -1371,41 +1371,41 @@ StackFrames ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.238.clone.clone (param_0.1076: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1076 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1014 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1076), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.441 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1014), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} -} - -%fused_computation.318.clone.clone (param_0.1077: f32[4,128], param_1.1243: bf16[4,128,4096], param_2.1068: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1068 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.379 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1068), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1243 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1016 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1243), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1077 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1661 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1077), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1660 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1016, %mul.1661), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1015 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1660), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.378 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.379, %convert_element_type.1015), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.371 (param_0.1091: f32[4096,128256], param_1.1254: f32[4,128], param_2.1090: bf16[4,128,4096], param_3.784: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { - %param_1.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1090 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.784 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.230.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1254, %param_2.1090, %param_3.784), kind=kLoop, calls=%fused_computation.318.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1091 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %fusion.211.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1091), kind=kLoop, calls=%fused_computation.238.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.81.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.230.clone.1, %fusion.211.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} +%fused_computation.237.clone.clone (param_0.1090: f32[4096,128256]) -> bf16[4096,128256,1] { + %param_0.1090 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %convert_element_type.1026 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + ROOT %bitcast.447 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1026), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +} + +%fused_computation.317.clone.clone (param_0.1091: f32[4,128], param_1.1257: bf16[4,128,4096], param_2.1077: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1077 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1077), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1257 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1028 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1257), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1091 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1595 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1091), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1594 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1028, %mul.1595), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1027 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1594), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1027), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.371 (param_0.1105: f32[4096,128256], param_1.1268: f32[4,128], param_2.1099: bf16[4,128,4096], param_3.788: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { + %param_1.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1099 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.788 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.240.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1268, %param_2.1099, %param_3.788), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1105 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %fusion.221.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1105), kind=kLoop, calls=%fused_computation.237.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %convolution.87.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.240.clone.1, %fusion.221.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} %constant.992 = bf16[]{:T(256)} constant(-inf) - %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.81.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} - ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.81.clone.1) + %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.87.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.87.clone.1) } -%fused_computation.372 (param_0.1088: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.1088 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %bitcast.444 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1088), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.444) +%fused_computation.372 (param_0.1102: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { + %param_0.1102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %bitcast.450 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1102), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.450) } %convert_element_type.525.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { @@ -1414,13 +1414,13 @@ StackFrames ROOT %add.624 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.121.clone.clone (param_0.1229: bf16[4,4096], param_1.1363: s32[]) -> bf16[4096] { - %param_0.1229 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1363 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.121.clone.clone (param_0.1242: bf16[4,4096], param_1.1376: s32[]) -> bf16[4096] { + %param_0.1242 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) %constant.1116 = s32[]{:T(128)} constant(0) - %dynamic_slice.310 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1229, %param_1.1363, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %dynamic_slice.316 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1242, %param_1.1376, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} %constant.1117 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.310, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.316, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %region_12.14 (reduce_sum.108: f32[], reduce_sum.109: f32[]) -> f32[] { @@ -1429,70 +1429,70 @@ StackFrames ROOT %reduce_sum.113 = f32[]{:T(128)} add(%reduce_sum.108, %reduce_sum.109), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.58.clone.clone (param_0.1230: bf16[4,4,128,4096], param_1.1364: s32[]) -> f32[4,128] { - %param_0.1230 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1364 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.58.clone.clone (param_0.1243: bf16[4,4,128,4096], param_1.1377: s32[]) -> f32[4,128] { + %param_0.1243 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1377 = s32[]{:T(128)S(6)} parameter(1) %constant.1118 = s32[]{:T(128)} constant(0) - %dynamic_slice.311 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1230, %param_1.1364, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.543 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.311), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1081 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.543), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.142 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1081, %convert_element_type.1081), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %dynamic_slice.317 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1243, %param_1.1377, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.548 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1093 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.548), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.214 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1093, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} %constant.1119 = f32[]{:T(128)} constant(0) - ROOT %reduce.175 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.142, %constant.1119), dimensions={2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + ROOT %reduce.175 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.214, %constant.1119), dimensions={2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} } -%fused_computation.143.clone.1.clone (param_0.1231: f32[4,128]) -> f32[4,128] { - %param_0.1231 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.143.clone.1.clone (param_0.1244: f32[4,128]) -> f32[4,128] { + %param_0.1244 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1121 = f32[]{:T(128)} constant(0.000244140625) %closed_call.81 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1121), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1231, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1244, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1120 = f32[]{:T(128)} constant(1e-05) %closed_call.80 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1120), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.858 = f32[4,128]{1,0:T(4,128)} add(%div.842, %closed_call.80), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.97 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.858), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.24.clone.1.clone.clone (param_0.1245: bf16[4,4096,32,128], param_1.1374: s32[]) -> bf16[4096,32,128,1] { - %param_0.1245 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1374 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.24.clone.1.clone.clone (param_0.1258: bf16[4,4096,32,128], param_1.1387: s32[]) -> bf16[4096,32,128,1] { + %param_0.1258 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) %constant.1134 = s32[]{:T(128)} constant(0) - %dynamic_slice.317 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1245, %param_1.1374, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.554 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.323 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1258, %param_1.1387, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.559 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.91.clone.clone (param_0.1246: f32[4,128], param_1.1375: bf16[4,4,128,4096], param_2.1167: s32[], param_3.843: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.843 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.843), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1375 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1167 = s32[]{:T(128)S(6)} parameter(2) +%fused_computation.91.clone.clone (param_0.1259: f32[4,128], param_1.1388: bf16[4,4,128,4096], param_2.1176: s32[], param_3.847: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.847 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.847), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1388 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1176 = s32[]{:T(128)S(6)} parameter(2) %constant.1135 = s32[]{:T(128)} constant(0) - %dynamic_slice.318 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1375, %param_2.1167, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.556 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.318), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1089 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.556), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1246 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1775 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1246), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1774 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1089, %mul.1775), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1088 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1774), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1088), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.555 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.36.clone.clone (param_0.1247: bf16[4,4096,32,128], param_1.1376: s32[], param_2.1168: f32[4,128], param_3.844: bf16[4,4,128,4096], param_4.525: bf16[4096]) -> bf16[4,128,32,128] { - %param_2.1168 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.844 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) - %param_4.525 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.332 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1168, %param_3.844, %param_1.1376, %param_4.525), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1247 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.331 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1247, %param_1.1376), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.107 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.332, %fusion.331), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.70.clone.clone (param_0.1248: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { - %param_0.1248 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1248), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %dynamic_slice.324 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1388, %param_2.1176, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.561 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1101 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1259 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1709 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1259), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1708 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1101, %mul.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1100 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1708), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1100), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.560 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.427), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.36.clone.clone (param_0.1260: bf16[4,4096,32,128], param_1.1389: s32[], param_2.1177: f32[4,128], param_3.848: bf16[4,4,128,4096], param_4.530: bf16[4096]) -> bf16[4,128,32,128] { + %param_2.1177 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.848 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1389 = s32[]{:T(128)S(6)} parameter(1) + %param_4.530 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.343 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1177, %param_3.848, %param_1.1389, %param_4.530), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1260 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.342 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1260, %param_1.1389), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.113 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.343, %fusion.342), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.70.clone.clone (param_0.1261: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { + %param_0.1261 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.129 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1248), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.187 = (bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) } @@ -1502,172 +1502,172 @@ StackFrames %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} %constant.1123 = s32[]{:T(128)} constant(2) %closed_call.83 = s32[64]{0:T(128)} broadcast(%constant.1123), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1765 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1082 = f32[64]{0:T(128)} convert(%mul.1765), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %mul.1699 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1094 = f32[64]{0:T(128)} convert(%mul.1699), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %constant.1122 = f32[]{:T(128)} constant(0.0078125) %closed_call.82 = f32[64]{0:T(128)} broadcast(%constant.1122), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1082, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1094, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.84, %div.843), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} } -%fused_computation.117.clone.clone (param_0.1232: f32[64], param_1.1365: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1365 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1365), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1232 = f32[64]{0:T(128)S(1)} parameter(0) - %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1232), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} +%fused_computation.117.clone.clone (param_0.1245: f32[64], param_1.1378: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1378 = f32[4,128]{1,0:T(4,128)} parameter(1) + %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1378), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %param_0.1245 = f32[64]{0:T(128)S(1)} parameter(0) + %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1245), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %div.844 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.846, %div.845), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1083 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1095 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} %convert_element_type.829.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1083, %convert_element_type.829.clone.3) + ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1095, %convert_element_type.829.clone.3) } -%fused_computation.120.clone.clone (param_0.1239: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1239 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.120.clone.clone (param_0.1252: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1252 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1130 = bf16[]{:T(256)} constant(-inf) - %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1239, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1239, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.61, %pad.60), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.554 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.119.clone.clone (param_0.1233: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1233 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.119.clone.clone (param_0.1246: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1246 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1125 = bf16[]{:T(256)} constant(-inf) - %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1233, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1233, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.44 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.59, %pad.58), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.544 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.73.clone.clone (param_0.1249: bf16[4,128,32,64], param_1.1377: bf16[4,128,32,64], param_2.1169: bf16[4,128,32,128], param_3.845: bf16[4,128,128], param_4.526: bf16[4,128,128]) -> bf16[4,32,128,128] { - %param_2.1169 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.526 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1779 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.526), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1777 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1169, %mul.1779), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1377 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.73.clone.clone (param_0.1262: bf16[4,128,32,64], param_1.1390: bf16[4,128,32,64], param_2.1178: bf16[4,128,32,128], param_3.849: bf16[4,128,128], param_4.531: bf16[4,128,128]) -> bf16[4,32,128,128] { + %param_2.1178 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.531 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1713 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.531), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1711 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1178, %mul.1713), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1390 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1136 = bf16[]{:T(256)} constant(-inf) - %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1377, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1249 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1249, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1390, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1262 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1262, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.47 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.65, %pad.64), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.845 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1778 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.845), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1776 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1778), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1777, %mul.1776), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.557 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.90.clone.clone (param_0.1241: f32[4,128], param_1.1371: bf16[4,4,128,4096], param_2.1164: s32[], param_3.840: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.840 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.422 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.840), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1371 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1164 = s32[]{:T(128)S(6)} parameter(2) + %param_3.849 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1712 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.849), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1710 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1712), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1711, %mul.1710), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.562 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.90.clone.clone (param_0.1254: f32[4,128], param_1.1384: bf16[4,4,128,4096], param_2.1173: s32[], param_3.844: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.844 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.844), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1384 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1173 = s32[]{:T(128)S(6)} parameter(2) %constant.1132 = s32[]{:T(128)} constant(0) - %dynamic_slice.316 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1371, %param_2.1164, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.316), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1087 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1241 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1769 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1241), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1768 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1087, %mul.1769), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1086 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1768), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.421 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.422, %convert_element_type.1086), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.64.clone.1.clone.clone (param_0.1240: bf16[4,4096,8,128], param_1.1370: s32[]) -> bf16[4096,8,128,1] { - %param_0.1240 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1370 = s32[]{:T(128)S(6)} parameter(1) + %dynamic_slice.322 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1384, %param_2.1173, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.557 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1099 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1703 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1254), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1702 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1099, %mul.1703), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1098 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1098), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.556 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.425), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.64.clone.1.clone.clone (param_0.1253: bf16[4,4096,8,128], param_1.1383: s32[]) -> bf16[4096,8,128,1] { + %param_0.1253 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1383 = s32[]{:T(128)S(6)} parameter(1) %constant.1131 = s32[]{:T(128)} constant(0) - %dynamic_slice.315 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1240, %param_1.1370, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.550 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.315), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.321 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1383, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.555 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.89.clone.clone (param_0.1242: bf16[4,4096,8,128], param_1.1372: s32[], param_2.1165: f32[4,128], param_3.841: bf16[4,4,128,4096], param_4.523: bf16[4096]) -> bf16[4,128,8,128] { - %param_2.1165 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.841 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1372 = s32[]{:T(128)S(6)} parameter(1) - %param_4.523 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.329 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1165, %param_3.841, %param_1.1372, %param_4.523), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1242 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.330 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1242, %param_1.1372), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.329, %fusion.330), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.89.clone.clone (param_0.1255: bf16[4,4096,8,128], param_1.1385: s32[], param_2.1174: f32[4,128], param_3.845: bf16[4,4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,8,128] { + %param_2.1174 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.845 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) + %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.340 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1174, %param_3.845, %param_1.1385, %param_4.528), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1255 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.341 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1255, %param_1.1385), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.112 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.340, %fusion.341), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } -%fused_computation.106.clone.clone (param_0.1243: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1243 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1243), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.106.clone.clone (param_0.1256: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { + %param_0.1256 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1243), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.186 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) } -%fused_computation.109.clone.clone (param_0.1244: bf16[4,128,8,64], param_1.1373: bf16[4,128,8,64], param_2.1166: bf16[4,128,8,128], param_3.842: bf16[4,128,128], param_4.524: bf16[4,128,128]) -> bf16[4,8,128,128] { - %param_2.1166 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.524 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1773 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.524), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1771 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1166, %mul.1773), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1373 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.109.clone.clone (param_0.1257: bf16[4,128,8,64], param_1.1386: bf16[4,128,8,64], param_2.1175: bf16[4,128,8,128], param_3.846: bf16[4,128,128], param_4.529: bf16[4,128,128]) -> bf16[4,8,128,128] { + %param_2.1175 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.529 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1707 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.529), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1705 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1175, %mul.1707), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1386 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1133 = bf16[]{:T(256)} constant(-inf) - %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1373, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1244 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1244, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1386, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1257 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1257, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.46 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.63, %pad.62), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.842 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1772 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1770 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1772), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1771, %mul.1770), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_3.846 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1706 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.846), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1704 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1706), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1705, %mul.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.558 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.135.clone.clone (param_0.1235: bf16[4,4096,8,128], param_1.1367: s32[]) -> bf16[1,4096,8,128] { - %param_0.1235 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1367 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.135.clone.clone (param_0.1248: bf16[4,4096,8,128], param_1.1380: s32[]) -> bf16[1,4096,8,128] { + %param_0.1248 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) + %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) %constant.1128 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.313 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1235, %param_1.1367, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %dynamic_slice.319 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1248, %param_1.1380, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } -%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1236: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { - %param_0.1236 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1236), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.545 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1249: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { + %param_0.1249 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} + ROOT %bitcast.550 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.88.clone.clone.clone.clone (param_0.1237: f32[4,128], param_1.1368: bf16[4,4,128,4096], param_2.1162: s32[], param_3.838: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.838 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.420 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.838), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1368 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1162 = s32[]{:T(128)S(6)} parameter(2) +%fused_computation.88.clone.clone.clone.clone (param_0.1250: f32[4,128], param_1.1381: bf16[4,4,128,4096], param_2.1171: s32[], param_3.842: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.842 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1381 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) %constant.1129 = s32[]{:T(128)} constant(0) - %dynamic_slice.314 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1368, %param_2.1162, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.547 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.314), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1085 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.547), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1237 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1767 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1237), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1766 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1085, %mul.1767), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1084 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1766), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.419 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.420, %convert_element_type.1084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.546 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.419), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.114.clone.clone (param_0.1238: bf16[1,4096,8,128], param_1.1369: f32[4,128], param_2.1163: bf16[4,4,128,4096], param_3.839: s32[], param_4.522: bf16[4096]) -> bf16[4,8,128,128] { - %param_1.1369 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1163 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.839 = s32[]{:T(128)S(6)} parameter(3) - %param_4.522 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.328 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1369, %param_2.1163, %param_3.839, %param_4.522), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1238 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.327 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1238), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.105 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.328, %fusion.327), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.548 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.146.clone.clone (param_0.1273: f32[4,32,128,128]) -> (f32[4,32,128], f32[4,32,128,1]) { - %param_0.1273 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1273), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.570 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.192 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.570, %slice.11) + %dynamic_slice.320 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1381, %param_2.1171, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1097 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1250 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1701 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1250), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1700 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1097, %mul.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1096 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1700), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1096), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.114.clone.clone (param_0.1251: bf16[1,4096,8,128], param_1.1382: f32[4,128], param_2.1172: bf16[4,4,128,4096], param_3.843: s32[], param_4.527: bf16[4096]) -> bf16[4,8,128,128] { + %param_1.1382 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1172 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.843 = s32[]{:T(128)S(6)} parameter(3) + %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.339 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1382, %param_2.1172, %param_3.843, %param_4.527), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1251 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.338 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1251), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.111 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.339, %fusion.338), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.366.clone.clone (param_0.1286: f32[4,32,128,128]) -> (f32[4,32,128,1], f32[4,32,128]) { + %param_0.1286 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) + %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1286), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.262.clone.3 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} + ROOT %tuple.192 = (f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}, f32[4,32,128]{2,1,0:T(8,128)S(1)}) tuple(%slice.11, %bitcast.262.clone.3) } %region_13.16 (reduce_sum.120: f32[], reduce_sum.121: f32[]) -> f32[] { @@ -1676,36 +1676,36 @@ StackFrames ROOT %reduce_sum.122 = f32[]{:T(128)} add(%reduce_sum.120, %reduce_sum.121), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1250: bf16[4,32,128,4096], param_1.1378: s32[]) -> bf16[32,128,4096,1] { - %param_0.1250 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1378 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1263: bf16[4,32,128,4096], param_1.1391: s32[]) -> bf16[32,128,4096,1] { + %param_0.1263 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) %constant.1137 = s32[]{:T(128)} constant(0) - %dynamic_slice.319 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1250, %param_1.1378, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.558 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.319), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1263, %param_1.1391, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.563 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1251: bf16[4,32,128,128]) -> bf16[4,128,32,128] { - %param_0.1251 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.559 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1264: bf16[4,32,128,128]) -> bf16[4,128,32,128] { + %param_0.1264 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.564 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.61.clone.clone (param_0.1252: bf16[4,32,128,4096], param_1.1379: s32[], param_2.1170: bf16[4,32,128,128], param_3.846: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { - %param_3.846 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.61.clone.clone (param_0.1265: bf16[4,32,128,4096], param_1.1392: s32[], param_2.1179: bf16[4,32,128,128], param_3.850: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { + %param_3.850 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1392 = s32[]{:T(128)S(6)} parameter(1) %constant.365.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.210.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.846, %param_1.1379, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.210.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1170 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.80.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1170), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1252 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.79.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1252, %param_1.1379), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.60.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.80.clone.3, %fusion.79.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %dynamic_slice.208.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.850, %param_1.1392, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.208.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %param_2.1179 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.83.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1179), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_0.1265 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.82.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1265, %param_1.1392), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.62.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.83.clone.3, %fusion.82.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %add.635.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.207.clone.3, %bitcast.182.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1090 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.143 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1090, %convert_element_type.1090), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %convert_element_type.1102 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.215 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1102, %convert_element_type.1102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} %constant.1138 = f32[]{:T(128)} constant(0) - %reduce.177 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.143, %constant.1138), dimensions={2}, to_apply=%region_13.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %reduce.177 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.215, %constant.1138), dimensions={2}, to_apply=%region_13.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.188 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.177, %add.635.clone.3) } @@ -1715,140 +1715,140 @@ StackFrames ROOT %add.623 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.122.clone.clone (param_0.1234: bf16[4,4096], param_1.1366: s32[]) -> bf16[4096] { - %param_0.1234 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1366 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.122.clone.clone (param_0.1247: bf16[4,4096], param_1.1379: s32[]) -> bf16[4096] { + %param_0.1247 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) %constant.1126 = s32[]{:T(128)} constant(0) - %dynamic_slice.312 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1234, %param_1.1366, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %dynamic_slice.318 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1247, %param_1.1379, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} %constant.1127 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.312, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.318, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.12.clone.clone.clone (param_0.1253: bf16[4,14336,4096], param_1.1380: s32[]) -> bf16[14336,4096,1] { - %param_0.1253 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.12.clone.clone.clone (param_0.1266: bf16[4,14336,4096], param_1.1393: s32[]) -> bf16[14336,4096,1] { + %param_0.1266 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1393 = s32[]{:T(128)S(6)} parameter(1) %constant.1139 = s32[]{:T(128)} constant(0) - %dynamic_slice.320 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1380, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.561 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.326 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1266, %param_1.1393, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.566 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.326), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %bitcast_fusion.3.clone.clone (bitcast_input.12: bf16[4,128,4096]) -> bf16[4,128,4096] { %bitcast_input.12 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.560 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) + ROOT %bitcast.565 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) } -%fused_computation.13.clone.clone (param_0.1254: bf16[4,128,4096], param_1.1381: bf16[4,14336,4096], param_2.1171: s32[]) -> bf16[14336,4,128] { - %param_1.1381 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) - %fusion.333 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1381, %param_2.1171), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1254 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %fusion.334 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1254), kind=kLoop, calls=%bitcast_fusion.3.clone.clone - ROOT %convolution.108 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.333, %fusion.334), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.13.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1394: bf16[4,14336,4096], param_2.1180: s32[]) -> bf16[14336,4,128] { + %param_1.1394 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) + %fusion.344 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1394, %param_2.1180), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.345 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1267), kind=kLoop, calls=%bitcast_fusion.3.clone.clone + ROOT %convolution.114 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.344, %fusion.345), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} } -%fused_computation.144.clone.1.clone (param_0.1255: f32[4,128]) -> f32[4,128] { - %param_0.1255 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.144.clone.1.clone (param_0.1268: f32[4,128]) -> f32[4,128] { + %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1141 = f32[]{:T(128)} constant(0.000244140625) %closed_call.86 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1141), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1255, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1140 = f32[]{:T(128)} constant(1e-05) %closed_call.85 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1140), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.861 = f32[4,128]{1,0:T(4,128)} add(%div.847, %closed_call.85), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.98 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.11.clone.1.clone.clone (param_0.1259: bf16[4,4096,14336], param_1.1385: s32[]) -> bf16[4096,14336,1] { - %param_0.1259 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.11.clone.1.clone.clone (param_0.1272: bf16[4,4096,14336], param_1.1398: s32[]) -> bf16[4096,14336,1] { + %param_0.1272 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1398 = s32[]{:T(128)S(6)} parameter(1) %constant.1143 = s32[]{:T(128)} constant(0) - %dynamic_slice.322 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1259, %param_1.1385, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.563 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.328 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1272, %param_1.1398, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.568 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.96.clone.2.clone.clone (param_0.1260: f32[4,128], param_1.1386: bf16[4,128,4096], param_2.1174: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1174), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1094 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1260 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1783 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1260), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1782 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1094, %mul.1783), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1093 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1782), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.96.clone.2.clone.clone (param_0.1273: f32[4,128], param_1.1399: bf16[4,128,4096], param_2.1183: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1183 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.432 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1183), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1399 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1106 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1717 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1273), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1716 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1106, %mul.1717), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1105 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1716), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.431 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.432, %convert_element_type.1105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.23.clone.clone (param_0.1261: bf16[4,4096,14336], param_1.1387: s32[], param_2.1175: f32[4,128], param_3.848: bf16[4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1175 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.848 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.338 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1175, %param_3.848, %param_4.528), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1261 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) - %fusion.337 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1261, %param_1.1387), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.110 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.338, %fusion.337), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.23.clone.clone (param_0.1274: bf16[4,4096,14336], param_1.1400: s32[], param_2.1184: f32[4,128], param_3.852: bf16[4,128,4096], param_4.533: bf16[4096]) -> bf16[4,128,14336] { + %param_2.1184 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.533 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.349 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1184, %param_3.852, %param_4.533), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1274 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1400 = s32[]{:T(128)S(6)} parameter(1) + %fusion.348 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1274, %param_1.1400), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.116 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.349, %fusion.348), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } -%fused_computation.14.clone.1.clone.clone (param_0.1262: bf16[4,4096,14336], param_1.1388: s32[]) -> bf16[4096,14336,1] { - %param_0.1262 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1388 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.14.clone.1.clone.clone (param_0.1275: bf16[4,4096,14336], param_1.1401: s32[]) -> bf16[4096,14336,1] { + %param_0.1275 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1401 = s32[]{:T(128)S(6)} parameter(1) %constant.1144 = s32[]{:T(128)} constant(0) - %dynamic_slice.323 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1262, %param_1.1388, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.564 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.329 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1275, %param_1.1401, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.569 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.329), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.39.clone.1.clone.clone (param_0.1263: bf16[14336,4,128], param_1.1389: bf16[4,128,14336]) -> bf16[4,128,14336] { - %param_1.1389 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.39.clone.1.clone.clone (param_0.1276: bf16[14336,4,128], param_1.1402: bf16[4,128,14336]) -> bf16[4,128,14336] { + %param_1.1402 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1145 = bf16[]{:T(256)} constant(1) %jit_silu_.44 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1145), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1402), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.69 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.862 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.848 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.862), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1785 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1389, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %param_0.1263 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) - %bitcast.565 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - ROOT %mul.1784 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1785, %bitcast.565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1719 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1402, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %param_0.1276 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.570 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + ROOT %mul.1718 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1719, %bitcast.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} } -%fused_computation.21.clone.clone (param_0.1264: bf16[4,4096,14336], param_1.1390: s32[], param_2.1176: bf16[14336,4,128], param_3.849: bf16[4,128,14336]) -> bf16[4,128,4096] { - %param_2.1176 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %param_3.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1176, %param_3.849), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_0.1264 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1390 = s32[]{:T(128)S(6)} parameter(1) - %fusion.339 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1264, %param_1.1390), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.111 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.339), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.21.clone.clone (param_0.1277: bf16[4,4096,14336], param_1.1403: s32[], param_2.1185: bf16[14336,4,128], param_3.853: bf16[4,128,14336]) -> bf16[4,128,4096] { + %param_2.1185 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) + %param_3.853 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1185, %param_3.853), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %param_0.1277 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1403 = s32[]{:T(128)S(6)} parameter(1) + %fusion.350 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1277, %param_1.1403), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.350), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} } -%fused_computation.14.clone.clone.clone (param_0.1256: bf16[4,4096,14336], param_1.1382: s32[]) -> bf16[4096,14336,1] { - %param_0.1256 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1382 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.14.clone.clone.clone (param_0.1269: bf16[4,4096,14336], param_1.1395: s32[]) -> bf16[4096,14336,1] { + %param_0.1269 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1395 = s32[]{:T(128)S(6)} parameter(1) %constant.1142 = s32[]{:T(128)} constant(0) - %dynamic_slice.321 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1256, %param_1.1382, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.562 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.327 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1269, %param_1.1395, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.567 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.96.clone.1.clone.clone (param_0.1257: f32[4,128], param_1.1383: bf16[4,128,4096], param_2.1172: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1172 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1172), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1092 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1257 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1781 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1257), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1780 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1092, %mul.1781), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1091 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1780), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1091), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.96.clone.1.clone.clone (param_0.1270: f32[4,128], param_1.1396: bf16[4,128,4096], param_2.1181: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1181 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1181), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1396 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1104 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1270 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1715 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1270), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1714 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1104, %mul.1715), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1103 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1714), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.430, %convert_element_type.1103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.20.clone.clone (param_0.1258: bf16[4,4096,14336], param_1.1384: s32[], param_2.1173: f32[4,128], param_3.847: bf16[4,128,4096], param_4.527: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1173 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.847 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.336 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1173, %param_3.847, %param_4.527), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1258 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1384 = s32[]{:T(128)S(6)} parameter(1) - %fusion.335 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1258, %param_1.1384), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.109 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.336, %fusion.335), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.20.clone.clone (param_0.1271: bf16[4,4096,14336], param_1.1397: s32[], param_2.1182: f32[4,128], param_3.851: bf16[4,128,4096], param_4.532: bf16[4096]) -> bf16[4,128,14336] { + %param_2.1182 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.532 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.347 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1182, %param_3.851, %param_4.532), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1271 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1397 = s32[]{:T(128)S(6)} parameter(1) + %fusion.346 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1271, %param_1.1397), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.115 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.347, %fusion.346), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } %region_14.17 (reduce_sum.126: f32[], reduce_sum.127: f32[]) -> f32[] { @@ -1857,63 +1857,63 @@ StackFrames ROOT %reduce_sum.128 = f32[]{:T(128)} add(%reduce_sum.126, %reduce_sum.127), metadata={op_name="checkpoint/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1265: bf16[4,4096,14336], param_1.1391: s32[]) -> bf16[4096,14336,1] { - %param_0.1265 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1278: bf16[4,4096,14336], param_1.1404: s32[]) -> bf16[4096,14336,1] { + %param_0.1278 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1404 = s32[]{:T(128)S(6)} parameter(1) %constant.1146 = s32[]{:T(128)} constant(0) - %dynamic_slice.324 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1265, %param_1.1391, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.566 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.330 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1278, %param_1.1404, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.571 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1266: bf16[4,128,14336], param_1.1392: bf16[4,128,14336], param_2.1177: bf16[14336,4,128]) -> bf16[4,128,14336] { - %param_2.1177 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %bitcast.567 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1177), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %param_1.1392 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %mul.1790 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.567, %param_1.1392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} +%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1279: bf16[4,128,14336], param_1.1405: bf16[4,128,14336], param_2.1186: bf16[14336,4,128]) -> bf16[4,128,14336] { + %param_2.1186 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.572 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1186), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + %param_1.1405 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %mul.1724 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.572, %param_1.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1147 = bf16[]{:T(256)} constant(1) %jit_silu_.45 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1147), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %param_0.1266 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1266), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %param_0.1279 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.70 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.131), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.863 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.70, %jit_silu_.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.45, %add.863), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1789 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1790, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1788 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1266, %mul.1790), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1723 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1724, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1722 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1279, %mul.1724), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} %sub.98 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} subtract(%jit_silu_.45, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/sub" stack_frame_id=0} - %mul.1787 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1786 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1788, %mul.1787), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1789, %mul.1786), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} -} - -%fused_computation.63.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1393: bf16[4096], param_2.1178: bf16[4,128,4096], param_3.850: bf16[4,4096,14336], param_4.529: s32[], param_5.425: bf16[4,128,14336], param_6.291: bf16[4,128,14336], param_7.188: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { - %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1096 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1267), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_2.1178 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_5.425 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) - %param_6.291 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) - %param_7.188 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) - %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.425, %param_6.291, %param_7.188), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} - %param_3.850 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.529 = s32[]{:T(128)S(6)} parameter(4) - %fusion.91.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.850, %param_4.529), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.64.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.91.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1178, %convolution.64.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1393 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) - %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1393), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.430), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1095 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.429), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1791 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1096, %convert_element_type.1095), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1721 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1720 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1722, %mul.1721), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1723, %mul.1720), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} +} + +%fused_computation.63.clone.clone (param_0.1280: bf16[4,128,4096], param_1.1406: bf16[4096], param_2.1187: bf16[4,128,4096], param_3.854: bf16[4,4096,14336], param_4.534: s32[], param_5.435: bf16[4,128,14336], param_6.304: bf16[4,128,14336], param_7.200: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { + %param_0.1280 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1108 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_2.1187 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_5.435 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) + %param_6.304 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) + %param_7.200 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) + %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.435, %param_6.304, %param_7.200), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} + %param_3.854 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) + %param_4.534 = s32[]{:T(128)S(6)} parameter(4) + %fusion.79.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.854, %param_4.534), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.60.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.79.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1187, %convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %param_1.1406 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %dot_general.434 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1406), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.433 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.434), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1107 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.433), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} + %mul.1725 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1108, %convert_element_type.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1148 = f32[]{:T(128)} constant(0) - %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1791, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} + %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1725, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.189 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.178, %add_any.132.clone.3) } -%fused_computation.140.clone.clone (param_0.1268: f32[4,128], param_1.1394: f32[4,128]) -> f32[4,128] { - %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1394 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.140.clone.clone (param_0.1281: f32[4,128], param_1.1407: f32[4,128]) -> f32[4,128] { + %param_0.1281 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1407 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.1152 = f32[]{:T(128)} constant(0.000244140625) %closed_call.89 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1152), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1394, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1407, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1151 = f32[]{:T(128)} constant(1e-05) %closed_call.88 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1151), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.864 = f32[4,128]{1,0:T(4,128)} add(%div.851, %closed_call.88), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} @@ -1921,11 +1921,11 @@ StackFrames %div.850 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.99, %add.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1150 = f32[]{:T(128)} constant(-0.5) %closed_call.87 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1150), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1794 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1793 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %mul.1794), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1728 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1727 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1281, %mul.1728), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1149 = f32[]{:T(128)} constant(0.00048828125) - %mul.1795 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - ROOT %mul.1792 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1793, %mul.1795), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1729 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + ROOT %mul.1726 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1727, %mul.1729), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} } %region_20.24 (dot_general.187: bf16[], dot_general.188: bf16[]) -> bf16[] { @@ -1934,29 +1934,29 @@ StackFrames ROOT %add.173 = bf16[]{:T(256)} add(%dot_general.187, %dot_general.188), metadata={op_name="add.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.94.clone.clone (param_0.1269: bf16[4,128,4096], param_1.1395: f32[4,128], param_2.1179: bf16[4,128,4096], param_3.851: bf16[4,128,4096], param_4.530: f32[4,128], param_5.426: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { - %param_0.1269 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_2.1179 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %convert_element_type.1098 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1179), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_1.1395 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1797 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1395), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1796 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1098, %mul.1797), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1097 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1796), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1269, %convert_element_type.1097), metadata={op_name="multiply.204"} +%fused_computation.94.clone.clone (param_0.1282: bf16[4,128,4096], param_1.1408: f32[4,128], param_2.1188: bf16[4,128,4096], param_3.855: bf16[4,128,4096], param_4.535: f32[4,128], param_5.436: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { + %param_0.1282 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_2.1188 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1110 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1188), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_1.1408 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1731 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1408), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1730 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1109 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1730), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %convert_element_type.1109), metadata={op_name="multiply.204"} %constant.1153 = bf16[]{:T(256)} constant(0) %reduce.179 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.271, %constant.1153), dimensions={0,1}, to_apply=%region_20.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_5.426 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) - %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.426), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1269, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.855 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_5.436 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.436), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} %convert_element_type.753.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.263.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1178.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1797), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_4.530 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %mul.1187.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.530), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %mul.1177.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1098, %mul.1187.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %add_any.126.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1178.clone.3, %mul.1177.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %mul.1142.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %param_4.535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %mul.1151.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.535), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1141.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1151.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %add_any.126.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1142.clone.3, %mul.1141.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} %convert_element_type.751.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.126.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.851, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.855, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} ROOT %tuple.190 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.179, %add_any.124.clone.3) } @@ -1966,35 +1966,35 @@ StackFrames ROOT %add.169 = f32[]{:T(128)} add(%dot_general.184, %dot_general.185), metadata={op_name="add.31"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1270: bf16[4,32,128,4096], param_1.1396: s32[]) -> bf16[32,128,4096,1] { - %param_0.1270 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1396 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1283: bf16[4,32,128,4096], param_1.1409: s32[]) -> bf16[32,128,4096,1] { + %param_0.1283 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1409 = s32[]{:T(128)S(6)} parameter(1) %constant.1154 = s32[]{:T(128)} constant(0) - %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1270, %param_1.1396, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.568 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.331 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1283, %param_1.1409, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.573 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1271: bf16[4,128,4096]) -> bf16[4,128,4096,1] { - %param_0.1271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.569 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1271), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} +%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1284: bf16[4,128,4096]) -> bf16[4,128,4096,1] { + %param_0.1284 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.574 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} } -%fused_computation.66.clone.clone (param_0.1272: bf16[4,32,128,128], param_1.1397: bf16[4,32,128,4096], param_2.1180: s32[], param_3.852: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { - %param_0.1272 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1272), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} - %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %fusion.84.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.852), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1397 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) - %fusion.83.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1397, %param_2.1180), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.62.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.84.clone.3, %fusion.83.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.66.clone.clone (param_0.1285: bf16[4,32,128,128], param_1.1410: bf16[4,32,128,4096], param_2.1189: s32[], param_3.856: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { + %param_0.1285 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} + %param_3.856 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %fusion.95.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.856), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %param_1.1410 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1189 = s32[]{:T(128)S(6)} parameter(2) + %fusion.94.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1410, %param_2.1189), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.64.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.95.clone.3, %fusion.94.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} %constant.619.clone.3 = bf16[]{:T(256)} constant(0.25) %div.442.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.619.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} - %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.62.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} + %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.64.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} %bitcast.209.clone.3 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%div.441.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} %convert.123 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.209.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert.1" stack_frame_id=0} %multiply.272 = f32[4,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.124, %convert.123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/multiply" stack_frame_id=0} %constant.1155 = f32[]{:T(128)} constant(0) - %dot_general.431 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} - ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.431, %bitcast.209.clone.3) + %dot_general.435 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} + ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.435, %bitcast.209.clone.3) } diff --git a/tests/utils/reference_hlo_qwen3_1.7b.txt b/tests/utils/reference_hlo_qwen3_1.7b.txt index 4004648b77..f1ede66966 100644 --- a/tests/utils/reference_hlo_qwen3_1.7b.txt +++ b/tests/utils/reference_hlo_qwen3_1.7b.txt @@ -32,7 +32,7 @@ StackFrames %param_1.5 = s32[512]{0:T(512)S(1)} parameter(1) %reshape.451 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.5), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} %transpose.466 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.451), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)} parameter(2) + %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)S(1)} parameter(2) %reshape.452 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} reshape(%param_2.4), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} %transpose.467 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} transpose(%reshape.452), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} ROOT %scatter.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.466, %transpose.467), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_42.47.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} @@ -50,43 +50,43 @@ StackFrames ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.386, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.277 (param_0.1367: f32[151936,2048], param_1.1549: f32[], param_2.1311: f32[], param_3.918: f32[], param_4.554: f32[151936,2048], param_5.467: f32[], param_6.356: bf16[151936,2048], param_7.196: bf16[151936,2048,1], param_8.113: pred[], param_9.94: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { - %param_0.1367 = f32[151936,2048]{1,0:T(8,128)} parameter(0) +%fused_computation.277 (param_0.1368: f32[151936,2048], param_1.1556: f32[], param_2.1314: f32[], param_3.918: f32[], param_4.556: f32[151936,2048], param_5.468: f32[], param_6.358: bf16[151936,2048], param_7.201: bf16[151936,2048,1], param_8.118: pred[], param_9.97: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { + %param_0.1368 = f32[151936,2048]{1,0:T(8,128)} parameter(0) %param_3.918 = f32[]{:T(128)S(6)} parameter(3) - %mul.2002.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_3.918), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.113 = pred[]{:T(512)S(6)} parameter(8) - %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.113), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_7.196 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) - %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.196), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %mul.1926.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_3.918), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.118 = pred[]{:T(512)S(6)} parameter(8) + %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.118), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_7.201 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) + %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.201), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %convert_element_type.1409.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.464.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_6.356 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.356), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_6.358 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %add_any.197.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1409.clone.1, %convert_element_type.1408.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} - %param_5.467 = f32[]{:T(128)} parameter(5) - %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.467), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_5.468 = f32[]{:T(128)} parameter(5) + %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.859.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add_any.197.clone.1, %div.860.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.267.clone.1 = f32[151936,2048]{1,0:T(8,128)} select(%select_n.268.clone.1, %add_any.197.clone.1, %div.859.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1092.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.844.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1092.clone.1), dimensions={}, metadata={op_name="broadcast.74"} - %mul.2008.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_9.94 = f32[151936,2048]{1,0:T(8,128)} parameter(9) + %mul.1932.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_9.97 = f32[151936,2048]{1,0:T(8,128)} parameter(9) %constant.1096.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2009.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1096.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2007.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.94, %mul.2009.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.941.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.2008.clone.1, %mul.2007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1311 = f32[]{:T(128)S(6)} parameter(2) - %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1311), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1933.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1096.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1931.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.97, %mul.1933.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.941.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1932.clone.1, %mul.1931.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) + %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.65.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %select_n.267.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1095.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2006.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1095.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2004.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %mul.2006.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.554 = f32[151936,2048]{1,0:T(8,128)} parameter(4) + %mul.1930.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1095.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1928.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %mul.1930.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.556 = f32[151936,2048]{1,0:T(8,128)} parameter(4) %constant.1094.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2005.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1094.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2003.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.554, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.940.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.2004.clone.1, %mul.2003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1549 = f32[]{:T(128)S(6)} parameter(1) - %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1549), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1929.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1094.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1927.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.556, %mul.1929.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.940.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1928.clone.1, %mul.1927.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) + %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.854.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.940.clone.1, %div.855.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.62.clone.1 = f32[151936,2048]{1,0:T(8,128)} sqrt(%div.854.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1093.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -94,14 +94,14 @@ StackFrames %add.938.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%sqrt.62.clone.1, %add.939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.426.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%div.856.clone.1, %add.938.clone.1), metadata={op_name="multiply.61"} %div.853.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.941.clone.1, %multiply.426.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2001.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1367, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.937.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%div.853.clone.1, %mul.2001.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2000.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%mul.2002.clone.1, %add.937.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1367, %mul.2000.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.175 = f32[151936,2048]{1,0:T(8,128)} multiply(%add.936.clone.1, %add.936.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1198 = f32[]{:T(128)} constant(0) - %reduce.176 = f32[]{:T(128)} reduce(%square.175, %constant.1198), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1198), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1925.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1368, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.937.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%div.853.clone.1, %mul.1925.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1924.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%mul.1926.clone.1, %add.937.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1368, %mul.1924.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.214 = f32[151936,2048]{1,0:T(8,128)} multiply(%add.936.clone.1, %add.936.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1200 = f32[]{:T(128)} constant(0) + %reduce.176 = f32[]{:T(128)} reduce(%square.214, %constant.1200), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.144 = (f32[]{:T(128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.176, %add.936.clone.1, %add.940.clone.1, %add.941.clone.1, %reduce.178.clone.1) } @@ -111,64 +111,64 @@ StackFrames ROOT %reduce_sum.319 = f32[]{:T(128)} add(%reduce_sum.317, %reduce_sum.318), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.367.clone.clone (param_0.1354: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1287: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1287 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1287), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1445 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1354 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2151 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1354), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2150 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1445, %mul.2151), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1444 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2150), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.478 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.479, %convert_element_type.1444), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.367.clone.clone (param_0.1355: f32[4,128], param_1.1549: bf16[4,128,2048], param_2.1290: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1290 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.480 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1290), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1549 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1451 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1355 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2083 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1355), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.2082 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1451, %mul.2083), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.480, %convert_element_type.1450), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.289.clone.clone.clone (param_0.1355: bf16[4,128,151936], param_1.1543: s32[4,128], param_2.1288: f32[4,128], param_3.911: f32[4,128], param_4.544: bf16[4,128], param_5.445: f32[4,128]) -> bf16[4,128,151936] { - %param_5.445 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2155 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.445), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.289.clone.clone.clone (param_0.1356: bf16[4,128,151936], param_1.1550: s32[4,128], param_2.1291: f32[4,128], param_3.911: f32[4,128], param_4.546: bf16[4,128], param_5.446: f32[4,128]) -> bf16[4,128,151936] { + %param_5.446 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.2087 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.446), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2154 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1355 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1448 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1355), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.544 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.544), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1448, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %mul.2086 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1356 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1454 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1356), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.546 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1454, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.62 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2153 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2154, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2153, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1543), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.2085 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2086, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1291 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1291), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2085, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1550 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.48 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.47 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1447 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1447), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2152 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2155, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1446 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2152), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.281 (param_0.1380: bf16[151936,2048], param_1.1562: f32[4,128], param_2.1324: bf16[4,128,2048], param_3.931: bf16[2048], param_4.567: bf16[4,128,151936], param_5.480: s32[4,128], param_6.369: f32[4,128], param_7.209: f32[4,128], param_8.126: bf16[4,128], param_9.95: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { - %param_4.567 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) - %param_5.480 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.369 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.209 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) - %param_8.126 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) - %param_9.95 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.567, %param_5.480, %param_6.369, %param_7.209, %param_8.126, /*index=5*/%param_9.95), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1562 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1324 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1453 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1453), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2084 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2087, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1452 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.281 (param_0.1381: bf16[151936,2048], param_1.1569: f32[4,128], param_2.1327: bf16[4,128,2048], param_3.931: bf16[2048], param_4.569: bf16[4,128,151936], param_5.481: s32[4,128], param_6.371: f32[4,128], param_7.214: f32[4,128], param_8.131: bf16[4,128], param_9.98: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { + %param_4.569 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) + %param_5.481 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.371 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.214 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) + %param_8.131 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) + %param_9.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.569, %param_5.481, %param_6.371, %param_7.214, %param_8.131, /*index=5*/%param_9.98), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1327 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) %param_3.931 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.268.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1562, %param_2.1324, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.268.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1569, %param_2.1327, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.269.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %bitcast.333 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%convolution.86.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %convert_element_type.1323 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_0.1380 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1380), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_0.1381 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %add_any.184 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1323, %convert_element_type.1322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} - %mul.1800 = f32[151936,2048]{1,0:T(8,128)} multiply(%add_any.184, %add_any.184), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1211 = f32[]{:T(128)} constant(0) - %reduce.177 = f32[]{:T(128)} reduce(%mul.1800, %constant.1211), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %square.215 = f32[151936,2048]{1,0:T(8,128)} multiply(%add_any.184, %add_any.184), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1213 = f32[]{:T(128)} constant(0) + %reduce.177 = f32[]{:T(128)} reduce(%square.215, %constant.1213), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.166 = (f32[]{:T(128)}, bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.177, %convolution.86.clone.1) } @@ -178,23 +178,23 @@ StackFrames ROOT %reduce_sum.394 = f32[]{:T(128)} add(%reduce_sum.389, %reduce_sum.393), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.288 (param_0.1391: bf16[4,128,151936], param_1.1570: f32[4,128], param_2.1327: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { - %param_2.1327 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1327), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.288 (param_0.1392: bf16[4,128,151936], param_1.1577: f32[4,128], param_2.1330: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { + %param_2.1330 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1330), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.25 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.24 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1391 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} %param_3.933 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) %sub.73 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.933), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.64 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1340, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1570 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1570), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1577 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1577), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.60 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.1223 = f32[]{:T(128)} constant(0) - %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.109"} - %mul.1805 = f32[4,128,151936]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.769), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1805, %constant.1223), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %constant.1225 = f32[]{:T(128)} constant(0) + %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1225), dimensions={}, metadata={op_name="broadcast.109"} + %mul.1765 = f32[4,128,151936]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.769), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1765, %constant.1225), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_9.12 (reduce_sum.186: f32[], reduce_sum.190: f32[]) -> f32[] { @@ -203,15 +203,15 @@ StackFrames ROOT %reduce_sum.191 = f32[]{:T(128)} add(%reduce_sum.186, %reduce_sum.190), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.293 (param_0.1392: bf16[4,128,151936], param_1.1571: bf16[4,128]) -> f32[4,128] { - %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1571 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1571), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.293 (param_0.1393: bf16[4,128,151936], param_1.1578: bf16[4,128]) -> f32[4,128] { + %param_0.1393 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1578 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1578), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.70 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1346, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.54 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.1224 = f32[]{:T(128)} constant(0) - ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1224), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %constant.1226 = f32[]{:T(128)} constant(0) + ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1226), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_33.38 (reduce_sum.269: f32[], reduce_sum.270: f32[]) -> f32[] { @@ -220,12 +220,12 @@ StackFrames ROOT %reduce_sum.274 = f32[]{:T(128)} add(%reduce_sum.269, %reduce_sum.270), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.298 (param_0.1386: f32[4,6144,2048]) -> f32[] { - %param_0.1386 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) - %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1810 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%bitcast.347, %bitcast.347), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1217 = f32[]{:T(128)} constant(0) - ROOT %reduce.181 = f32[]{:T(128)} reduce(%mul.1810, %constant.1217), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.298 (param_0.1387: f32[4,6144,2048]) -> f32[] { + %param_0.1387 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) + %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.218 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%bitcast.347, %bitcast.347), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1219 = f32[]{:T(128)} constant(0) + ROOT %reduce.181 = f32[]{:T(128)} reduce(%square.218, %constant.1219), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_32.37 (reduce_sum.263: f32[], reduce_sum.267: f32[]) -> f32[] { @@ -240,35 +240,35 @@ StackFrames ROOT %reduce_sum.262 = f32[]{:T(128)} add(%reduce_sum.260, %reduce_sum.261), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.300 (param_0.1387: f32[4,2048,6144], param_1.1566: f32[4,2048,6144]) -> (f32[], f32[]) { - %param_0.1387 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) - %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1813 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.351, %bitcast.351), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1218 = f32[]{:T(128)} constant(0) - %reduce.182 = f32[]{:T(128)} reduce(%mul.1813, %constant.1218), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1566 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) - %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1566), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1816.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.355.clone.1, %bitcast.355.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.183.clone.1 = f32[]{:T(128)} reduce(%mul.1816.clone.1, %constant.1218), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.300 (param_0.1388: f32[4,2048,6144], param_1.1573: f32[4,2048,6144]) -> (f32[], f32[]) { + %param_0.1388 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) + %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.221 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.351, %bitcast.351), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1220 = f32[]{:T(128)} constant(0) + %reduce.182 = f32[]{:T(128)} reduce(%square.221, %constant.1220), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1573 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) + %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.224.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.355.clone.1, %bitcast.355.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.183.clone.1 = f32[]{:T(128)} reduce(%square.224.clone.1, %constant.1220), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.167 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.182, %reduce.183.clone.1) } -%fused_computation.303 (param_0.900: f32[6144,4,2048]) -> bf16[4,6144,2048] { - %param_0.900 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) - %copy.192 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.900), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.303 (param_0.901: f32[6144,4,2048]) -> bf16[4,6144,2048] { + %param_0.901 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) + %copy.190 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.901), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} + ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.190), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.304 (param_0.902: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.902 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.193 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.902), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.304 (param_0.903: f32[2048,4,6144]) -> bf16[4,2048,6144] { + %param_0.903 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.191 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.903), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} + ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.305 (param_0.904: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.904 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.194 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.904), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.305 (param_0.905: f32[2048,4,6144]) -> bf16[4,2048,6144] { + %param_0.905 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.192 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.905), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} + ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_62.67 (reduce_sum.416: f32[], reduce_sum.417: f32[]) -> f32[] { @@ -283,39 +283,39 @@ StackFrames ROOT %reduce_sum.340 = f32[]{:T(128)} add(%reduce_sum.338, %reduce_sum.339), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.306 (param_0.1376: f32[6144,4,2048], param_1.1558: f32[], param_2.1320: f32[], param_3.927: f32[], param_4.563: f32[6144,4,2048], param_5.476: f32[], param_6.365: f32[4,6144,2048], param_7.205: pred[], param_8.122: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { - %param_0.1376 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) +%fused_computation.306 (param_0.1377: f32[6144,4,2048], param_1.1565: f32[], param_2.1323: f32[], param_3.927: f32[], param_4.565: f32[6144,4,2048], param_5.477: f32[], param_6.367: f32[4,6144,2048], param_7.210: pred[], param_8.127: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { + %param_0.1377 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) %param_3.927 = f32[]{:T(128)S(6)} parameter(3) - %mul.2074.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_3.927), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.205 = pred[]{:T(512)S(6)} parameter(7) - %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.365 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) - %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.476 = f32[]{:T(128)} parameter(5) - %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1998.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_3.927), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.210 = pred[]{:T(512)S(6)} parameter(7) + %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.210), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.367 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) + %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.477 = f32[]{:T(128)} parameter(5) + %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.931.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%bitcast.482.clone.1, %div.932.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.303.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} select(%select_n.304.clone.1, %bitcast.482.clone.1, %div.931.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1146.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.886.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1146.clone.1), dimensions={}, metadata={op_name="broadcast.83"} - %mul.2080.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.122 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) + %mul.2004.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.127 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) %constant.1150.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2081.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1150.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2079.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.122, %mul.2081.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.989.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2080.clone.1, %mul.2079.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) - %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2005.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1150.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2003.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.127, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.989.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2004.clone.1, %mul.2003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) + %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.74.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %select_n.303.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1149.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2078.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1149.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2076.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%integer_pow.74.clone.1, %mul.2078.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.563 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) + %mul.2002.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1149.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2000.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%integer_pow.74.clone.1, %mul.2002.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.565 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) %constant.1148.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2077.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1148.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2075.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.563, %mul.2077.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.988.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2076.clone.1, %mul.2075.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) - %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2001.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1148.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1999.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.565, %mul.2001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.988.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2000.clone.1, %mul.1999.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1565 = f32[]{:T(128)S(6)} parameter(1) + %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1565), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.926.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.988.clone.1, %div.927.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.71.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} sqrt(%div.926.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1147.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -323,14 +323,14 @@ StackFrames %add.986.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%sqrt.71.clone.1, %add.987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.435.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%div.928.clone.1, %add.986.clone.1), metadata={op_name="multiply.52"} %div.925.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.989.clone.1, %multiply.435.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2073.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1376, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.985.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%div.925.clone.1, %mul.2073.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2072.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%mul.2074.clone.1, %add.985.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1376, %mul.2072.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.176 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%add.984.clone.1, %add.984.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1207 = f32[]{:T(128)} constant(0) - %reduce.184 = f32[]{:T(128)} reduce(%square.176, %constant.1207), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1207), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1997.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.985.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%div.925.clone.1, %mul.1997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1996.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%mul.1998.clone.1, %add.985.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1377, %mul.1996.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.225 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%add.984.clone.1, %add.984.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1209 = f32[]{:T(128)} constant(0) + %reduce.184 = f32[]{:T(128)} reduce(%square.225, %constant.1209), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.145 = (f32[]{:T(128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.184, %add.984.clone.1, %add.988.clone.1, %add.989.clone.1, %reduce.187.clone.1) } @@ -346,39 +346,39 @@ StackFrames ROOT %reduce_sum.337 = f32[]{:T(128)} add(%reduce_sum.332, %reduce_sum.333), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.307 (param_0.1377: f32[2048,4,6144], param_1.1559: f32[], param_2.1321: f32[], param_3.928: f32[], param_4.564: f32[2048,4,6144], param_5.477: f32[], param_6.366: f32[4,2048,6144], param_7.206: pred[], param_8.123: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1377 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) +%fused_computation.307 (param_0.1378: f32[2048,4,6144], param_1.1566: f32[], param_2.1324: f32[], param_3.928: f32[], param_4.566: f32[2048,4,6144], param_5.478: f32[], param_6.368: f32[4,2048,6144], param_7.211: pred[], param_8.128: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { + %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) %param_3.928 = f32[]{:T(128)S(6)} parameter(3) - %mul.2084.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.928), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.206 = pred[]{:T(512)S(6)} parameter(7) - %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.366 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.477 = f32[]{:T(128)} parameter(5) - %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2008.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.928), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.211 = pred[]{:T(512)S(6)} parameter(7) + %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.211), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.368 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.478 = f32[]{:T(128)} parameter(5) + %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.939.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.484.clone.1, %div.940.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.307.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.308.clone.1, %bitcast.484.clone.1, %div.939.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1152.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.892.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1152.clone.1), dimensions={}, metadata={op_name="broadcast.85"} - %mul.2088.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.123 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) + %mul.2012.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.128 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) %constant.1156.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.891.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1156.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2087.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.123, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.994.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2088.clone.1, %mul.2087.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) - %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2011.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.128, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.994.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2012.clone.1, %mul.2011.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1324 = f32[]{:T(128)S(6)} parameter(2) + %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1324), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.75.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %select_n.307.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1155.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.890.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1155.clone.1), dimensions={}, metadata={op_name="broadcast.73"} - %mul.2086.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.75.clone.1, %broadcast.890.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.564 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) + %mul.2010.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.75.clone.1, %broadcast.890.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.566 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) %constant.1154.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.889.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1154.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2085.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.564, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.993.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2086.clone.1, %mul.2085.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) - %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2009.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.566, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.993.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2010.clone.1, %mul.2009.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1566 = f32[]{:T(128)S(6)} parameter(1) + %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1566), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.934.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.993.clone.1, %div.935.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.72.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.934.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1153.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -386,14 +386,14 @@ StackFrames %add.992.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.72.clone.1, %broadcast.887.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.436.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.936.clone.1, %add.992.clone.1), metadata={op_name="multiply.51"} %div.933.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.994.clone.1, %multiply.436.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2083.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.991.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.933.clone.1, %mul.2083.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2082.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2084.clone.1, %add.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1377, %mul.2082.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.177 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.990.clone.1, %add.990.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1208 = f32[]{:T(128)} constant(0) - %reduce.185 = f32[]{:T(128)} reduce(%square.177, %constant.1208), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1208), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2007.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.991.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.933.clone.1, %mul.2007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2006.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2008.clone.1, %add.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2006.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.226 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.990.clone.1, %add.990.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1210 = f32[]{:T(128)} constant(0) + %reduce.185 = f32[]{:T(128)} reduce(%square.226, %constant.1210), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1210), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.146 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.185, %add.990.clone.1, %add.993.clone.1, %add.994.clone.1, %reduce.188.clone.1) } @@ -409,39 +409,39 @@ StackFrames ROOT %reduce_sum.331 = f32[]{:T(128)} add(%reduce_sum.326, %reduce_sum.330), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.308 (param_0.1378: f32[2048,4,6144], param_1.1560: f32[], param_2.1322: f32[], param_3.929: f32[], param_4.565: f32[2048,4,6144], param_5.478: f32[], param_6.367: f32[4,2048,6144], param_7.207: pred[], param_8.124: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) +%fused_computation.308 (param_0.1379: f32[2048,4,6144], param_1.1567: f32[], param_2.1325: f32[], param_3.929: f32[], param_4.567: f32[2048,4,6144], param_5.479: f32[], param_6.369: f32[4,2048,6144], param_7.212: pred[], param_8.129: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { + %param_0.1379 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) %param_3.929 = f32[]{:T(128)S(6)} parameter(3) - %mul.2091.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.929), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.207 = pred[]{:T(512)S(6)} parameter(7) - %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.367 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.478 = f32[]{:T(128)} parameter(5) - %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2015.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.929), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.212 = pred[]{:T(512)S(6)} parameter(7) + %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.212), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.369 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.369), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.479 = f32[]{:T(128)} parameter(5) + %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.947.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.486.clone.1, %div.948.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.311.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.312.clone.1, %bitcast.486.clone.1, %div.947.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1158.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.898.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1158.clone.1), dimensions={}, metadata={op_name="broadcast.85"} - %mul.2095.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.124 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) + %mul.2019.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.129 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) %constant.1162.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.897.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1162.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2094.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.124, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.999.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2095.clone.1, %mul.2094.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) - %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2018.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.129, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.999.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2019.clone.1, %mul.2018.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1325 = f32[]{:T(128)S(6)} parameter(2) + %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1325), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.76.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %select_n.311.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1161.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.896.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1161.clone.1), dimensions={}, metadata={op_name="broadcast.73"} - %mul.2093.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.76.clone.1, %broadcast.896.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.565 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) + %mul.2017.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.76.clone.1, %broadcast.896.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.567 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) %constant.1160.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.895.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1160.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2092.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.565, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.998.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2093.clone.1, %mul.2092.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) - %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2016.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.567, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.998.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2017.clone.1, %mul.2016.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1567 = f32[]{:T(128)S(6)} parameter(1) + %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1567), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.942.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.998.clone.1, %div.943.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.73.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.942.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1159.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -449,14 +449,14 @@ StackFrames %add.997.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.73.clone.1, %broadcast.893.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.437.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.944.clone.1, %add.997.clone.1), metadata={op_name="multiply.50"} %div.941.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.999.clone.1, %multiply.437.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2090.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.996.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.941.clone.1, %mul.2090.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2089.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2091.clone.1, %add.996.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2089.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.178 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.995.clone.1, %add.995.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1209 = f32[]{:T(128)} constant(0) - %reduce.186 = f32[]{:T(128)} reduce(%square.178, %constant.1209), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2014.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1379, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.996.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.941.clone.1, %mul.2014.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2013.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2015.clone.1, %add.996.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1379, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.227 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.995.clone.1, %add.995.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1211 = f32[]{:T(128)} constant(0) + %reduce.186 = f32[]{:T(128)} reduce(%square.227, %constant.1211), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1211), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.147 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.186, %add.995.clone.1, %add.998.clone.1, %add.999.clone.1, %reduce.189.clone.1) } @@ -466,12 +466,12 @@ StackFrames ROOT %reduce_sum.304 = f32[]{:T(128)} add(%reduce_sum.302, %reduce_sum.303), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.324 (param_0.1381: f32[4,2048,16,128]) -> f32[] { - %param_0.1381 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1845 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%bitcast.362, %bitcast.362), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1212 = f32[]{:T(128)} constant(0) - ROOT %reduce.190 = f32[]{:T(128)} reduce(%mul.1845, %constant.1212), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.324 (param_0.1382: f32[4,2048,16,128]) -> f32[] { + %param_0.1382 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.230 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%bitcast.362, %bitcast.362), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1214 = f32[]{:T(128)} constant(0) + ROOT %reduce.190 = f32[]{:T(128)} reduce(%square.230, %constant.1214), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_38.43 (reduce_sum.296: f32[], reduce_sum.297: f32[]) -> f32[] { @@ -480,18 +480,18 @@ StackFrames ROOT %reduce_sum.298 = f32[]{:T(128)} add(%reduce_sum.296, %reduce_sum.297), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.326 (param_0.1382: f32[4,16,128,2048]) -> f32[] { - %param_0.1382 = f32[4,16,128,2048]{3,2,0,1:T(8,128)S(1)} parameter(0) - %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1848 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%bitcast.366, %bitcast.366), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1213 = f32[]{:T(128)} constant(0) - ROOT %reduce.191 = f32[]{:T(128)} reduce(%mul.1848, %constant.1213), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.326 (param_0.1383: f32[4,16,128,2048]) -> f32[] { + %param_0.1383 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.233 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%bitcast.366, %bitcast.366), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1215 = f32[]{:T(128)} constant(0) + ROOT %reduce.191 = f32[]{:T(128)} reduce(%square.233, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.327 (param_0.949: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { - %param_0.949 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) - %copy.195 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.949), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.195), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.327 (param_0.950: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { + %param_0.950 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) + %copy.193 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.950), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} + ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_68.73 (reduce_sum.449: f32[], reduce_sum.450: f32[]) -> f32[] { @@ -506,39 +506,39 @@ StackFrames ROOT %reduce_sum.373 = f32[]{:T(128)} add(%reduce_sum.368, %reduce_sum.372), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.328 (param_0.1370: f32[2048,4,16,128], param_1.1552: f32[], param_2.1314: f32[], param_3.921: f32[], param_4.557: f32[2048,4,16,128], param_5.470: f32[], param_6.359: f32[4,2048,16,128], param_7.199: pred[], param_8.116: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { - %param_0.1370 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.328 (param_0.1371: f32[2048,4,16,128], param_1.1559: f32[], param_2.1317: f32[], param_3.921: f32[], param_4.559: f32[2048,4,16,128], param_5.471: f32[], param_6.361: f32[4,2048,16,128], param_7.204: pred[], param_8.121: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { + %param_0.1371 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.921 = f32[]{:T(128)S(6)} parameter(3) - %mul.2026.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_3.921), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.199 = pred[]{:T(512)S(6)} parameter(7) - %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.199), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.359 = f32[4,2048,16,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.470 = f32[]{:T(128)} parameter(5) - %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1950.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_3.921), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.204 = pred[]{:T(512)S(6)} parameter(7) + %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.361 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.471 = f32[]{:T(128)} parameter(5) + %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.883.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%bitcast.470.clone.1, %div.884.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.279.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} select(%select_n.280.clone.1, %bitcast.470.clone.1, %div.883.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1110.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.858.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1110.clone.1), dimensions={}, metadata={op_name="broadcast.75"} - %mul.2032.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.116 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.121 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1114.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2033.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1114.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2031.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.116, %mul.2033.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.2032.clone.1, %mul.2031.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) - %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1114.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.121, %mul.1957.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1956.clone.1, %mul.1955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) + %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.68.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %select_n.279.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1113.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2030.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1113.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2028.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.68.clone.1, %mul.2030.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.557 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1113.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.68.clone.1, %mul.1954.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.559 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1112.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2029.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1112.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2027.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.2028.clone.1, %mul.2027.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1552 = f32[]{:T(128)S(6)} parameter(1) - %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1552), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1112.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1951.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.559, %mul.1953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1952.clone.1, %mul.1951.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) + %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.878.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.956.clone.1, %div.879.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.65.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} sqrt(%div.878.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1111.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -546,14 +546,14 @@ StackFrames %add.954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%sqrt.65.clone.1, %add.955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.429.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%div.880.clone.1, %add.954.clone.1), metadata={op_name="multiply.58"} %div.877.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.957.clone.1, %multiply.429.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2025.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1370, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%div.877.clone.1, %mul.2025.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2024.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%mul.2026.clone.1, %add.953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1370, %mul.2024.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.179 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%add.952.clone.1, %add.952.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1201 = f32[]{:T(128)} constant(0) - %reduce.192 = f32[]{:T(128)} reduce(%square.179, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1949.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%div.877.clone.1, %mul.1949.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1948.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%mul.1950.clone.1, %add.953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.1948.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.234 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%add.952.clone.1, %add.952.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1203 = f32[]{:T(128)} constant(0) + %reduce.192 = f32[]{:T(128)} reduce(%square.234, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.148 = (f32[]{:T(128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.192, %add.952.clone.1, %add.956.clone.1, %add.957.clone.1, %reduce.194.clone.1) } @@ -569,39 +569,39 @@ StackFrames ROOT %reduce_sum.367 = f32[]{:T(128)} add(%reduce_sum.365, %reduce_sum.366), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.329 (param_0.1371: f32[16,4,128,2048], param_1.1553: f32[], param_2.1315: f32[], param_3.922: f32[], param_4.558: f32[16,4,128,2048], param_5.471: f32[], param_6.360: f32[4,16,128,2048], param_7.200: pred[], param_8.117: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { - %param_0.1371 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.329 (param_0.1372: f32[16,4,128,2048], param_1.1560: f32[], param_2.1318: f32[], param_3.922: f32[], param_4.560: f32[16,4,128,2048], param_5.472: f32[], param_6.362: f32[4,16,128,2048], param_7.205: pred[], param_8.122: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { + %param_0.1372 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) %param_3.922 = f32[]{:T(128)S(6)} parameter(3) - %mul.2036.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_3.922), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.200 = pred[]{:T(512)S(6)} parameter(7) - %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.200), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.360 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.471 = f32[]{:T(128)} parameter(5) - %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_3.922), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.205 = pred[]{:T(512)S(6)} parameter(7) + %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.362 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.472 = f32[]{:T(128)} parameter(5) + %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.891.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%bitcast.472.clone.1, %div.892.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.283.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} select(%select_n.284.clone.1, %bitcast.472.clone.1, %div.891.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1116.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.860.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1116.clone.1), dimensions={}, metadata={op_name="broadcast.76"} - %mul.2042.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.117 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) + %mul.1966.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.122 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) %constant.1120.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2043.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1120.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2041.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.117, %mul.2043.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.2042.clone.1, %mul.2041.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) - %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1967.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1120.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1965.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.122, %mul.1967.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1966.clone.1, %mul.1965.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) + %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.69.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %select_n.283.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1119.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2040.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1119.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2038.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%integer_pow.69.clone.1, %mul.2040.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.558 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) + %mul.1964.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1119.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%integer_pow.69.clone.1, %mul.1964.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.560 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) %constant.1118.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2039.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1118.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2037.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.558, %mul.2039.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.2038.clone.1, %mul.2037.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1553 = f32[]{:T(128)S(6)} parameter(1) - %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1553), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1118.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.560, %mul.1963.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1962.clone.1, %mul.1961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) + %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.886.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.962.clone.1, %div.887.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.66.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} sqrt(%div.886.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1117.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -609,14 +609,14 @@ StackFrames %add.960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%sqrt.66.clone.1, %add.961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.430.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%div.888.clone.1, %add.960.clone.1), metadata={op_name="multiply.57"} %div.885.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.963.clone.1, %multiply.430.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2035.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%div.885.clone.1, %mul.2035.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2034.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%mul.2036.clone.1, %add.959.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.2034.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.180 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%add.958.clone.1, %add.958.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1202 = f32[]{:T(128)} constant(0) - %reduce.193 = f32[]{:T(128)} reduce(%square.180, %constant.1202), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1202), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1372, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%div.885.clone.1, %mul.1959.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%mul.1960.clone.1, %add.959.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1372, %mul.1958.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.235 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%add.958.clone.1, %add.958.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1204 = f32[]{:T(128)} constant(0) + %reduce.193 = f32[]{:T(128)} reduce(%square.235, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.149 = (f32[]{:T(128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.193, %add.958.clone.1, %add.962.clone.1, %add.963.clone.1, %reduce.195.clone.1) } @@ -632,23 +632,23 @@ StackFrames ROOT %reduce_sum.289 = f32[]{:T(128)} add(%reduce_sum.284, %reduce_sum.288), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.341 (param_0.1384: f32[4,2048,8,128], param_1.1564: f32[4,2048,8,128]) -> (f32[], f32[]) { - %param_0.1384 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) - %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1863 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1215 = f32[]{:T(128)} constant(0) - %reduce.196 = f32[]{:T(128)} reduce(%mul.1863, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1564 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) - %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1866.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.197.clone.1 = f32[]{:T(128)} reduce(%mul.1866.clone.1, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.341 (param_0.1385: f32[4,2048,8,128], param_1.1571: f32[4,2048,8,128]) -> (f32[], f32[]) { + %param_0.1385 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.238 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1217 = f32[]{:T(128)} constant(0) + %reduce.196 = f32[]{:T(128)} reduce(%square.238, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1571 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(1) + %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.241.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.197.clone.1 = f32[]{:T(128)} reduce(%square.241.clone.1, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.168 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.196, %reduce.197.clone.1) } -%fused_computation.344 (param_0.981: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.981 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %copy.196 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.981), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.196), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.344 (param_0.982: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { + %param_0.982 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.194 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.982), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} + ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_70.75 (reduce_sum.458: f32[], reduce_sum.459: f32[]) -> f32[] { @@ -663,39 +663,39 @@ StackFrames ROOT %reduce_sum.382 = f32[]{:T(128)} add(%reduce_sum.380, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.345 (param_0.1368: f32[2048,4,8,128], param_1.1550: f32[], param_2.1312: f32[], param_3.919: f32[], param_4.555: f32[2048,4,8,128], param_5.468: f32[], param_6.357: f32[4,2048,8,128], param_7.197: pred[], param_8.114: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1368 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.345 (param_0.1369: f32[2048,4,8,128], param_1.1557: f32[], param_2.1315: f32[], param_3.919: f32[], param_4.557: f32[2048,4,8,128], param_5.469: f32[], param_6.359: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { + %param_0.1369 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.919 = f32[]{:T(128)S(6)} parameter(3) - %mul.2012.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.919), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.197 = pred[]{:T(512)S(6)} parameter(7) - %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.357 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.357), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.468 = f32[]{:T(128)} parameter(5) - %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1936.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.919), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.202 = pred[]{:T(512)S(6)} parameter(7) + %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.359 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.469 = f32[]{:T(128)} parameter(5) + %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.867.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.466.clone.1, %div.868.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.271.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.272.clone.1, %bitcast.466.clone.1, %div.867.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1098.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.850.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1098.clone.1), dimensions={}, metadata={op_name="broadcast.80"} - %mul.2016.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.114 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1940.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1102.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.849.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1102.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.2015.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.114, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.946.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2016.clone.1, %mul.2015.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1312 = f32[]{:T(128)S(6)} parameter(2) - %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1312), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1939.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.946.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1940.clone.1, %mul.1939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) + %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.66.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %select_n.271.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1101.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.848.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1101.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.2014.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.66.clone.1, %broadcast.848.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.555 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1938.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.66.clone.1, %broadcast.848.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.557 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1100.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.847.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1100.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.2013.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.555, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.945.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2014.clone.1, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1550 = f32[]{:T(128)S(6)} parameter(1) - %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1937.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.945.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1938.clone.1, %mul.1937.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) + %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.862.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.945.clone.1, %div.863.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.63.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.862.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1099.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -703,14 +703,14 @@ StackFrames %add.944.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.845.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.427.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.864.clone.1, %add.944.clone.1), metadata={op_name="multiply.60"} %div.861.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.946.clone.1, %multiply.427.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2011.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1368, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.943.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.861.clone.1, %mul.2011.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2010.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.2012.clone.1, %add.943.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1368, %mul.2010.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.181 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.942.clone.1, %add.942.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1199 = f32[]{:T(128)} constant(0) - %reduce.198 = f32[]{:T(128)} reduce(%square.181, %constant.1199), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1199), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1935.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1369, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.943.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.861.clone.1, %mul.1935.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1934.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1936.clone.1, %add.943.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1369, %mul.1934.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.242 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.942.clone.1, %add.942.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1201 = f32[]{:T(128)} constant(0) + %reduce.198 = f32[]{:T(128)} reduce(%square.242, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.150 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.198, %add.942.clone.1, %add.945.clone.1, %add.946.clone.1, %reduce.200.clone.1) } @@ -726,39 +726,39 @@ StackFrames ROOT %reduce_sum.358 = f32[]{:T(128)} add(%reduce_sum.353, %reduce_sum.354), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.346 (param_0.1373: f32[2048,4,8,128], param_1.1555: f32[], param_2.1317: f32[], param_3.924: f32[], param_4.560: f32[2048,4,8,128], param_5.473: f32[], param_6.362: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1373 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) +%fused_computation.346 (param_0.1374: f32[2048,4,8,128], param_1.1562: f32[], param_2.1320: f32[], param_3.924: f32[], param_4.562: f32[2048,4,8,128], param_5.474: f32[], param_6.364: f32[4,2048,8,128], param_7.207: pred[], param_8.124: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { + %param_0.1374 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.924 = f32[]{:T(128)S(6)} parameter(3) - %mul.2053.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.924), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.202 = pred[]{:T(512)S(6)} parameter(7) - %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.362 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.473 = f32[]{:T(128)} parameter(5) - %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1977.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.924), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.207 = pred[]{:T(512)S(6)} parameter(7) + %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.364 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.474 = f32[]{:T(128)} parameter(5) + %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.907.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.476.clone.1, %div.908.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.291.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.292.clone.1, %bitcast.476.clone.1, %div.907.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1128.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.872.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1128.clone.1), dimensions={}, metadata={op_name="broadcast.80"} - %mul.2057.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1981.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.124 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1132.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.871.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1132.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.2056.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.973.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2057.clone.1, %mul.2056.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) - %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1980.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.124, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.973.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1981.clone.1, %mul.1980.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) + %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.71.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %select_n.291.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1131.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.870.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1131.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.2055.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.71.clone.1, %broadcast.870.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.560 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1979.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.71.clone.1, %broadcast.870.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.562 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1130.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.869.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1130.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.2054.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.560, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.972.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2055.clone.1, %mul.2054.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1555 = f32[]{:T(128)S(6)} parameter(1) - %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1555), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1978.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.562, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.972.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1979.clone.1, %mul.1978.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1562 = f32[]{:T(128)S(6)} parameter(1) + %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1562), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.902.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.972.clone.1, %div.903.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.68.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.902.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1129.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -766,32 +766,32 @@ StackFrames %add.971.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.68.clone.1, %broadcast.867.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.432.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.904.clone.1, %add.971.clone.1), metadata={op_name="multiply.55"} %div.901.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.973.clone.1, %multiply.432.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2052.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1373, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.970.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.901.clone.1, %mul.2052.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2051.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.2053.clone.1, %add.970.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1373, %mul.2051.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.182 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.969.clone.1, %add.969.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1204 = f32[]{:T(128)} constant(0) - %reduce.199 = f32[]{:T(128)} reduce(%square.182, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) + %mul.1976.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1374, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.970.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.901.clone.1, %mul.1976.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1975.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1977.clone.1, %add.970.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1374, %mul.1975.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.243 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.969.clone.1, %add.969.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1206 = f32[]{:T(128)} constant(0) + %reduce.199 = f32[]{:T(128)} reduce(%square.243, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) } -%fused_computation.362 (param_0.1055: bf16[4,128,2048], param_1.1114: f32[4,128], param_2.829: f32[4,128], param_3.497: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { - %param_3.497 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) +%fused_computation.362 (param_0.1056: bf16[4,128,2048], param_1.1117: f32[4,128], param_2.830: f32[4,128], param_3.495: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { + %param_3.495 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) %param_4.296 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.451 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.441 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.497, %dot_general.451), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.441), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.829 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1912 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.829), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1904 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1363, %mul.1912), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.1055 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1055), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1911 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1114), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1910 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1374, %mul.1911), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %add_any.193 = f32[4,128,2048]{2,1,0:T(8,128)} add(%mul.1904, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} + %dot_general.448 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.438 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.495, %dot_general.448), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.438), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_2.830 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %mul.1851 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.830), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1843 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1363, %mul.1851), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_0.1056 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1056), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1117 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1850 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1117), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1849 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1374, %mul.1850), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %add_any.193 = f32[4,128,2048]{2,1,0:T(8,128)} add(%mul.1843, %mul.1849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} ROOT %convert_element_type.1361 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.193), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} } @@ -801,12 +801,12 @@ StackFrames ROOT %reduce_sum.185 = f32[]{:T(128)} add(%reduce_sum.171, %reduce_sum.184), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.363 (param_0.1393: bf16[4,128,2048]) -> f32[4,128] { - %param_0.1393 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.185 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1365, %convert_element_type.1365), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} - %constant.1225 = f32[]{:T(128)} constant(0) - ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.185, %constant.1225), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} +%fused_computation.363 (param_0.1394: bf16[4,128,2048]) -> f32[4,128] { + %param_0.1394 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1394), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.246 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1365, %convert_element_type.1365), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} + %constant.1227 = f32[]{:T(128)} constant(0) + ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.246, %constant.1227), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_12.15 (reduce_sum.198: f32[], reduce_sum.199: f32[]) -> f32[] { @@ -815,17 +815,17 @@ StackFrames ROOT %reduce_sum.200 = f32[]{:T(128)} add(%reduce_sum.198, %reduce_sum.199), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.365 (param_0.1388: bf16[4,128,2048], param_1.1567: bf16[4,128,2048], param_2.1325: bf16[2048]) -> f32[4,128] { - %param_0.1388 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1388), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1325 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1325), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.440 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1567, %dot_general.450), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.440), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1908 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1372, %convert_element_type.1371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %constant.1219 = f32[]{:T(128)} constant(0) - ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1908, %constant.1219), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} +%fused_computation.365 (param_0.1389: bf16[4,128,2048], param_1.1574: bf16[4,128,2048], param_2.1328: bf16[2048]) -> f32[4,128] { + %param_0.1389 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1389), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1328 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.447 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1328), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1574, %dot_general.447), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.437), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.1847 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1372, %convert_element_type.1371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %constant.1221 = f32[]{:T(128)} constant(0) + ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1847, %constant.1221), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_10.13 (dot_general.190: bf16[], dot_general.191: bf16[]) -> bf16[] { @@ -834,64 +834,64 @@ StackFrames ROOT %add.419 = bf16[]{:T(256)} add(%dot_general.190, %dot_general.191), metadata={op_name="add.82"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.285.clone.clone (param_0.1350: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1350 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.530 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1350), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%fused_computation.285.clone.clone (param_0.1351: bf16[151936,2048]) -> bf16[151936,2048,1] { + %param_0.1351 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} } -%fused_computation.289.clone.1.clone.clone (param_0.1351: bf16[4,128,151936], param_1.1539: s32[4,128], param_2.1282: f32[4,128], param_3.906: f32[4,128], param_4.540: bf16[4,128], param_5.441: f32[4,128]) -> bf16[4,128,151936] { - %param_5.441 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2143 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.441), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.289.clone.1.clone.clone (param_0.1352: bf16[4,128,151936], param_1.1546: s32[4,128], param_2.1285: f32[4,128], param_3.906: f32[4,128], param_4.542: bf16[4,128], param_5.442: f32[4,128]) -> bf16[4,128,151936] { + %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.2075 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.442), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.906 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2142 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1351 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1438 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.540 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.540), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1438, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %mul.2074 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1352 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1444 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1352), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.542 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.542), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1444, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.60 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2141 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2142, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1282 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1282), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2141, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1539 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.2073 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2074, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1285 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1285), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2073, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1546 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.42 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.41 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1437 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1437), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2140 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2143, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1436 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2140), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convert_element_type.1443 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1443), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2072 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2075, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1442 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2072), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} } -%fused_computation.366 (param_0.1349: f32[4,128], param_1.1538: bf16[4,128,2048], param_2.1283: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.541: s32[4,128], param_5.442: f32[4,128], param_6.338: f32[4,128], param_7.194: bf16[4,128], param_8.111: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { +%fused_computation.366 (param_0.1350: f32[4,128], param_1.1545: bf16[4,128,2048], param_2.1286: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.543: s32[4,128], param_5.443: f32[4,128], param_6.340: f32[4,128], param_7.199: bf16[4,128], param_8.116: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { %param_3.907 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.541 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.338 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.194 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.111 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.541, %param_5.442, %param_6.338, %param_7.194, /*index=5*/%param_8.111), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1283 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) - %fusion.250.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1283), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.250.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %param_1.1538 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1349 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1923 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1349), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1922 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1384, %mul.1923), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1383 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1922), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_4.543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %param_5.443 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.340 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.199 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) + %param_8.116 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) + %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.543, %param_5.443, %param_6.340, %param_7.199, /*index=5*/%param_8.116), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_2.1286 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) + %fusion.251.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1286), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.251.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %param_1.1545 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1545), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1350 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1862 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1350), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1861 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1384, %mul.1862), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1383 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1861), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %multiply.420 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%convolution.84.clone.1, %convert_element_type.1383), metadata={op_name="multiply.362"} %constant.1050 = bf16[]{:T(256)} constant(0) %reduce.204 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%multiply.420, %constant.1050), dimensions={0,1}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} ROOT %tuple.165 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.204, %convolution.84.clone.1) } -%fused_computation.374 (param_0.1087: f32[64], param_1.1147: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1147 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1147), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.1087 = f32[64]{0:T(128)S(1)} parameter(0) - %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1087), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} +%fused_computation.374 (param_0.1088: f32[64], param_1.1150: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1150 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1150), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %param_0.1088 = f32[64]{0:T(128)S(1)} parameter(0) + %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1088), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %div.717 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.720, %div.718), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.717), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} %convert_element_type.1392 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} @@ -900,19 +900,19 @@ StackFrames ROOT %tuple.158 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1392, %convert_element_type.1391.clone.1) } -%fused_computation.375 (param_0.1084: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1084 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.375 (param_0.1085: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.1085 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1042 = bf16[]{:T(256)} constant(-inf) - %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1084, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1084, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.42 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.46, %pad.45), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } -%fused_computation.376 (param_0.1086: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1086 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.376 (param_0.1087: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.1087 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1041 = bf16[]{:T(256)} constant(-inf) - %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1086, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1086, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.43 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.48, %pad.47), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } @@ -928,16 +928,16 @@ StackFrames ROOT %reduce_sum.277 = f32[]{:T(128)} add(%reduce_sum.275, %reduce_sum.276), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.380 (param_0.1385: f32[4,2048], param_1.1565: f32[4,2048]) -> (f32[], f32[]) { - %param_0.1385 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1932 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.404, %bitcast.404), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1216 = f32[]{:T(128)} constant(0) - %reduce.205 = f32[]{:T(128)} reduce(%mul.1932, %constant.1216), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1565 = f32[4,2048]{1,0:T(4,128)} parameter(1) - %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1935.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.408.clone.1, %bitcast.408.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.206.clone.1 = f32[]{:T(128)} reduce(%mul.1935.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.380 (param_0.1386: f32[4,2048], param_1.1572: f32[4,2048]) -> (f32[], f32[]) { + %param_0.1386 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.249 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.404, %bitcast.404), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1218 = f32[]{:T(128)} constant(0) + %reduce.205 = f32[]{:T(128)} reduce(%square.249, %constant.1218), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1572 = f32[4,2048]{1,0:T(4,128)} parameter(1) + %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.252.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.408.clone.1, %bitcast.408.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.206.clone.1 = f32[]{:T(128)} reduce(%square.252.clone.1, %constant.1218), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.169 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.205, %reduce.206.clone.1) } @@ -953,39 +953,39 @@ StackFrames ROOT %reduce_sum.352 = f32[]{:T(128)} add(%reduce_sum.347, %reduce_sum.351), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.383 (param_0.1374: f32[2048,4], param_1.1556: f32[], param_2.1318: f32[], param_3.925: f32[], param_4.561: f32[2048,4], param_5.474: f32[], param_6.363: f32[4,2048], param_7.203: pred[], param_8.120: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1374 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.383 (param_0.1375: f32[2048,4], param_1.1563: f32[], param_2.1321: f32[], param_3.925: f32[], param_4.563: f32[2048,4], param_5.475: f32[], param_6.365: f32[4,2048], param_7.208: pred[], param_8.125: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.925 = f32[]{:T(128)S(6)} parameter(3) - %mul.2060.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.925), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.203 = pred[]{:T(512)S(6)} parameter(7) - %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.363 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.474 = f32[]{:T(128)} parameter(5) - %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1984.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.925), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.208 = pred[]{:T(512)S(6)} parameter(7) + %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.365 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.475 = f32[]{:T(128)} parameter(5) + %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.915.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.478.clone.1, %div.916.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.295.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.296.clone.1, %bitcast.478.clone.1, %div.915.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1134.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.878.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1134.clone.1), dimensions={}, metadata={op_name="broadcast.82"} - %mul.2064.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.120 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1988.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.125 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1138.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.877.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1138.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.2063.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.978.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2064.clone.1, %mul.2063.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) - %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1987.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.125, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.978.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1988.clone.1, %mul.1987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) + %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.72.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %select_n.295.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1137.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.876.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1137.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.2062.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.72.clone.1, %broadcast.876.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.561 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1986.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.72.clone.1, %broadcast.876.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.563 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1136.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.875.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1136.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.2061.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.977.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2062.clone.1, %mul.2061.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) - %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1985.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.563, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.977.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1986.clone.1, %mul.1985.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1563 = f32[]{:T(128)S(6)} parameter(1) + %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1563), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.910.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.977.clone.1, %div.911.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.69.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.910.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1135.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -993,14 +993,14 @@ StackFrames %add.976.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.69.clone.1, %broadcast.873.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.433.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.912.clone.1, %add.976.clone.1), metadata={op_name="multiply.54"} %div.909.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.978.clone.1, %multiply.433.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2059.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1374, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.975.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.909.clone.1, %mul.2059.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2058.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.2060.clone.1, %add.975.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1374, %mul.2058.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.186 = f32[2048,4]{0,1:T(4,128)} multiply(%add.974.clone.1, %add.974.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1205 = f32[]{:T(128)} constant(0) - %reduce.207 = f32[]{:T(128)} reduce(%square.186, %constant.1205), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1983.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.975.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.909.clone.1, %mul.1983.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1982.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1984.clone.1, %add.975.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.1982.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.253 = f32[2048,4]{0,1:T(4,128)} multiply(%add.974.clone.1, %add.974.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1207 = f32[]{:T(128)} constant(0) + %reduce.207 = f32[]{:T(128)} reduce(%square.253, %constant.1207), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1207), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.152 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.207, %add.974.clone.1, %add.977.clone.1, %add.978.clone.1, %reduce.209.clone.1) } @@ -1016,39 +1016,39 @@ StackFrames ROOT %reduce_sum.346 = f32[]{:T(128)} add(%reduce_sum.344, %reduce_sum.345), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.384 (param_0.1375: f32[2048,4], param_1.1557: f32[], param_2.1319: f32[], param_3.926: f32[], param_4.562: f32[2048,4], param_5.475: f32[], param_6.364: f32[4,2048], param_7.204: pred[], param_8.121: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.384 (param_0.1376: f32[2048,4], param_1.1564: f32[], param_2.1322: f32[], param_3.926: f32[], param_4.564: f32[2048,4], param_5.476: f32[], param_6.366: f32[4,2048], param_7.209: pred[], param_8.126: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1376 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.926 = f32[]{:T(128)S(6)} parameter(3) - %mul.2067.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.926), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.204 = pred[]{:T(512)S(6)} parameter(7) - %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.364 = f32[4,2048]{1,0:T(4,128)} parameter(6) - %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.475 = f32[]{:T(128)} parameter(5) - %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1991.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.926), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.209 = pred[]{:T(512)S(6)} parameter(7) + %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.209), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.366 = f32[4,2048]{1,0:T(4,128)} parameter(6) + %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.476 = f32[]{:T(128)} parameter(5) + %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.923.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.480.clone.1, %div.924.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.299.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.300.clone.1, %bitcast.480.clone.1, %div.923.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1140.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.884.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1140.clone.1), dimensions={}, metadata={op_name="broadcast.82"} - %mul.2071.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.121 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1995.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.126 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1144.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.883.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1144.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.2070.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.121, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.983.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2071.clone.1, %mul.2070.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) - %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1994.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.126, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.983.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1995.clone.1, %mul.1994.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) + %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.73.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %select_n.299.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1143.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.882.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1143.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.2069.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.73.clone.1, %broadcast.882.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.562 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1993.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.73.clone.1, %broadcast.882.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.564 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1142.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.881.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1142.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.2068.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.562, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.982.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2069.clone.1, %mul.2068.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) - %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1992.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.564, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.982.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1993.clone.1, %mul.1992.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1564 = f32[]{:T(128)S(6)} parameter(1) + %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1564), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.918.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.982.clone.1, %div.919.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.70.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.918.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1141.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1056,14 +1056,14 @@ StackFrames %add.981.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.70.clone.1, %broadcast.879.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.434.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.920.clone.1, %add.981.clone.1), metadata={op_name="multiply.53"} %div.917.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.983.clone.1, %multiply.434.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2066.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.980.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.917.clone.1, %mul.2066.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2065.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.2067.clone.1, %add.980.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.2065.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.187 = f32[2048,4]{0,1:T(4,128)} multiply(%add.979.clone.1, %add.979.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1206 = f32[]{:T(128)} constant(0) - %reduce.208 = f32[]{:T(128)} reduce(%square.187, %constant.1206), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1206), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1990.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1376, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.980.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.917.clone.1, %mul.1990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1989.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1991.clone.1, %add.980.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1376, %mul.1989.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.254 = f32[2048,4]{0,1:T(4,128)} multiply(%add.979.clone.1, %add.979.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1208 = f32[]{:T(128)} constant(0) + %reduce.208 = f32[]{:T(128)} reduce(%square.254, %constant.1208), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1208), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.153 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.208, %add.979.clone.1, %add.982.clone.1, %add.983.clone.1, %reduce.210.clone.1) } @@ -1073,12 +1073,12 @@ StackFrames ROOT %reduce_sum.197 = f32[]{:T(128)} add(%reduce_sum.192, %reduce_sum.193), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.395 (param_0.1389: bf16[2048]) -> f32[] { - %param_0.1389 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1952 = f32[2048]{0:T(1024)} multiply(%convert_element_type.1396, %convert_element_type.1396), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1220 = f32[]{:T(128)} constant(0) - ROOT %reduce.211 = f32[]{:T(128)} reduce(%mul.1952, %constant.1220), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.395 (param_0.1390: bf16[2048]) -> f32[] { + %param_0.1390 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.257 = f32[2048]{0:T(1024)} multiply(%convert_element_type.1396, %convert_element_type.1396), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1222 = f32[]{:T(128)} constant(0) + ROOT %reduce.211 = f32[]{:T(128)} reduce(%square.257, %constant.1222), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_59.64 (reduce_sum.401: f32[], reduce_sum.402: f32[]) -> f32[] { @@ -1093,39 +1093,39 @@ StackFrames ROOT %reduce_sum.325 = f32[]{:T(128)} add(%reduce_sum.323, %reduce_sum.324), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.396 (param_0.1379: f32[2048], param_1.1561: f32[], param_2.1323: f32[], param_3.930: f32[], param_4.566: f32[2048], param_5.479: f32[], param_6.368: bf16[2048], param_7.208: pred[], param_8.125: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { - %param_0.1379 = f32[2048]{0:T(1024)S(1)} parameter(0) +%fused_computation.396 (param_0.1380: f32[2048], param_1.1568: f32[], param_2.1326: f32[], param_3.930: f32[], param_4.568: f32[2048], param_5.480: f32[], param_6.370: bf16[2048], param_7.213: pred[], param_8.130: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { + %param_0.1380 = f32[2048]{0:T(1024)S(1)} parameter(0) %param_3.930 = f32[]{:T(128)S(6)} parameter(3) - %mul.2098.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_3.930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.208 = pred[]{:T(512)S(6)} parameter(7) - %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.368 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.479 = f32[]{:T(128)} parameter(5) - %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2022.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_3.930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.213 = pred[]{:T(512)S(6)} parameter(7) + %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.213), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.370 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_5.480 = f32[]{:T(128)} parameter(5) + %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.480), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.955.clone.1 = f32[2048]{0:T(1024)} divide(%convert_element_type.1411.clone.1, %div.956.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.315.clone.1 = f32[2048]{0:T(1024)} select(%select_n.316.clone.1, %convert_element_type.1411.clone.1, %div.955.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1164.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.900.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1164.clone.1), dimensions={}, metadata={op_name="broadcast.86"} - %mul.2104.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.125 = f32[2048]{0:T(1024)S(1)} parameter(8) + %mul.2028.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.130 = f32[2048]{0:T(1024)S(1)} parameter(8) %constant.1168.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2105.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1168.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2103.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.125, %mul.2105.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1005.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2104.clone.1, %mul.2103.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) - %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2029.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1168.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2027.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.130, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1005.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2028.clone.1, %mul.2027.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1326 = f32[]{:T(128)S(6)} parameter(2) + %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1326), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.77.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %select_n.315.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1167.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2102.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1167.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2100.clone.1 = f32[2048]{0:T(1024)} multiply(%integer_pow.77.clone.1, %mul.2102.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.566 = f32[2048]{0:T(1024)S(1)} parameter(4) + %mul.2026.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1167.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2024.clone.1 = f32[2048]{0:T(1024)} multiply(%integer_pow.77.clone.1, %mul.2026.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.568 = f32[2048]{0:T(1024)S(1)} parameter(4) %constant.1166.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2101.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1166.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2099.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.566, %mul.2101.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1004.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2100.clone.1, %mul.2099.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) - %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2025.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1166.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2023.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.568, %mul.2025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1004.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2024.clone.1, %mul.2023.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) + %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.950.clone.1 = f32[2048]{0:T(1024)} divide(%add.1004.clone.1, %div.951.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.74.clone.1 = f32[2048]{0:T(1024)} sqrt(%div.950.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1165.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1133,37 +1133,37 @@ StackFrames %add.1002.clone.1 = f32[2048]{0:T(1024)} add(%sqrt.74.clone.1, %add.1003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.438.clone.1 = f32[2048]{0:T(1024)} multiply(%div.952.clone.1, %add.1002.clone.1), metadata={op_name="multiply.49"} %div.949.clone.1 = f32[2048]{0:T(1024)} divide(%add.1005.clone.1, %multiply.438.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2097.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1379, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1001.clone.1 = f32[2048]{0:T(1024)} add(%div.949.clone.1, %mul.2097.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2096.clone.1 = f32[2048]{0:T(1024)} multiply(%mul.2098.clone.1, %add.1001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1379, %mul.2096.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.188 = f32[2048]{0:T(1024)} multiply(%add.1000.clone.1, %add.1000.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1210 = f32[]{:T(128)} constant(0) - %reduce.212 = f32[]{:T(128)} reduce(%square.188, %constant.1210), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1210), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.2021.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1380, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1001.clone.1 = f32[2048]{0:T(1024)} add(%div.949.clone.1, %mul.2021.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2020.clone.1 = f32[2048]{0:T(1024)} multiply(%mul.2022.clone.1, %add.1001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1380, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.258 = f32[2048]{0:T(1024)} multiply(%add.1000.clone.1, %add.1000.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1212 = f32[]{:T(128)} constant(0) + %reduce.212 = f32[]{:T(128)} reduce(%square.258, %constant.1212), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1212), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.156 = (f32[]{:T(128)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.212, %add.1000.clone.1, %add.1004.clone.1, %add.1005.clone.1, %reduce.213.clone.1) } -%fused_computation.402 (param_0.1149: s32[512]) -> s32[1024] { +%fused_computation.402 (param_0.1150: s32[512]) -> s32[1024] { %constant.972 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.815 = s32[1024]{0:T(1024)} broadcast(%constant.972), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.1149 = s32[512]{0:T(512)S(1)} parameter(0) + %param_0.1150 = s32[512]{0:T(512)S(1)} parameter(0) %constant.973 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1149, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1150, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %constant.971 = s32[] constant(151935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.814 = s32[1024]{0:T(1024)} broadcast(%constant.971), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.815, %pad.49, %broadcast.814), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } -%fused_computation.405 (param_0.1148: s32[4,128]) -> s32[512] { - %param_0.1148 = s32[4,128]{1,0:T(4,128)} parameter(0) +%fused_computation.405 (param_0.1149: s32[4,128]) -> s32[512] { + %param_0.1149 = s32[4,128]{1,0:T(4,128)} parameter(0) %constant.1065 = s32[]{:T(128)} constant(0) %broadcast.834 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1065), dimensions={}, metadata={op_name="broadcast.95"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1148, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} + %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1149, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} %constant.1051 = s32[]{:T(128)} constant(151936) %add.925 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1051), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1148, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1148), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} + %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1149, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} + %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1149), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} ROOT %bitcast.409 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} } @@ -1179,16 +1179,16 @@ StackFrames ROOT %reduce_sum.295 = f32[]{:T(128)} add(%reduce_sum.290, %reduce_sum.291), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.407 (param_0.1383: f32[4,128], param_1.1563: f32[4,128]) -> (f32[], f32[]) { - %param_0.1383 = f32[4,128]{1,0:T(4,128)} parameter(0) - %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1965 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.413, %bitcast.413), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %constant.1214 = f32[]{:T(128)} constant(0) - %reduce.214 = f32[]{:T(128)} reduce(%mul.1965, %constant.1214), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1563 = f32[4,128]{1,0:T(4,128)} parameter(1) - %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %mul.1968.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.417.clone.1, %bitcast.417.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.215.clone.1 = f32[]{:T(128)} reduce(%mul.1968.clone.1, %constant.1214), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.407 (param_0.1384: f32[4,128], param_1.1570: f32[4,128]) -> (f32[], f32[]) { + %param_0.1384 = f32[4,128]{1,0:T(4,128)} parameter(0) + %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.261 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.413, %bitcast.413), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1216 = f32[]{:T(128)} constant(0) + %reduce.214 = f32[]{:T(128)} reduce(%square.261, %constant.1216), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1570 = f32[4,128]{1,0:T(4,128)} parameter(1) + %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.264.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.417.clone.1, %bitcast.417.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.215.clone.1 = f32[]{:T(128)} reduce(%square.264.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.170 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.214, %reduce.215.clone.1) } @@ -1204,31 +1204,31 @@ StackFrames ROOT %reduce_sum.400 = f32[]{:T(128)} add(%reduce_sum.395, %reduce_sum.396), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.410 (param_0.1390: bf16[4,128], param_1.1569: f32[4,128], param_2.1326: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { +%fused_computation.410 (param_0.1391: bf16[4,128], param_1.1576: f32[4,128], param_2.1329: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { %param_3.932 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) %constant.1170.clone.1 = s32[]{:T(128)} constant(0) %broadcast.901.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1170.clone.1), dimensions={}, metadata={op_name="broadcast.95"} %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.932, %broadcast.901.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1569), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1390 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1390), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %param_1.1576 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1576), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_0.1391 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} %add.927 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %square.193 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %add.927), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} - %constant.1222 = f32[]{:T(128)} constant(0) - %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1222), dimensions={}, metadata={op_name="broadcast.99"} - %mul.1989 = f32[4,128]{1,0:T(4,128)} multiply(%square.193, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %mul.1969 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1989, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.216 = f32[]{:T(128)} reduce(%mul.1969, %constant.1222), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1326 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1326), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} - %add.904.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1989), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %mul.1970.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.904.clone.1, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1970.clone.1, %constant.1222), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %mul.1987.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %broadcast.831), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %square.269 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %add.927), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.1224 = f32[]{:T(128)} constant(0) + %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1224), dimensions={}, metadata={op_name="broadcast.99"} + %mul.1913 = f32[4,128]{1,0:T(4,128)} multiply(%square.269, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1893 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1913, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.216 = f32[]{:T(128)} reduce(%mul.1893, %constant.1224), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_2.1329 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1329), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.904.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1913), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1894.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.904.clone.1, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1894.clone.1, %constant.1224), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1911.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %broadcast.831), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %constant.1068.clone.1 = f32[]{:T(128)} constant(1) %add.922.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1068.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - %add.915.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1987.clone.1, %add.922.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.915.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1911.clone.1, %add.922.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}, pred[4,128]{1,0:T(4,128)(4,1)S(1)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%reduce.216, %reduce.219.clone.1, %ne.6.clone.1, %add.915.clone.1) } @@ -1244,39 +1244,39 @@ StackFrames ROOT %reduce_sum.379 = f32[]{:T(128)} add(%reduce_sum.374, %reduce_sum.375), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.411 (param_0.1369: f32[128,4], param_1.1551: f32[], param_2.1313: f32[], param_3.920: f32[], param_4.556: f32[128,4], param_5.469: f32[], param_6.358: f32[4,128], param_7.198: pred[], param_8.115: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1369 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.411 (param_0.1370: f32[128,4], param_1.1558: f32[], param_2.1316: f32[], param_3.920: f32[], param_4.558: f32[128,4], param_5.470: f32[], param_6.360: f32[4,128], param_7.203: pred[], param_8.120: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1370 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.920 = f32[]{:T(128)S(6)} parameter(3) - %mul.2019.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.920), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.198 = pred[]{:T(512)S(6)} parameter(7) - %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.198), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.358 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.469 = f32[]{:T(128)} parameter(5) - %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1943.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.920), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.203 = pred[]{:T(512)S(6)} parameter(7) + %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.360 = f32[4,128]{1,0:T(4,128)} parameter(6) + %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.470 = f32[]{:T(128)} parameter(5) + %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.875.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.468.clone.1, %div.876.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.275.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.276.clone.1, %bitcast.468.clone.1, %div.875.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1104.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.856.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1104.clone.1), dimensions={}, metadata={op_name="broadcast.78"} - %mul.2023.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.115 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1947.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.120 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1108.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.855.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1108.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.2022.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.115, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.951.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2023.clone.1, %mul.2022.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1313 = f32[]{:T(128)S(6)} parameter(2) - %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1313), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1946.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.951.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1947.clone.1, %mul.1946.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) + %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.67.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %select_n.275.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1107.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.854.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1107.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.2021.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.854.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.556 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1945.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.854.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.558 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1106.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.853.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1106.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.2020.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.556, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.950.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2021.clone.1, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1551 = f32[]{:T(128)S(6)} parameter(1) - %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1551), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1944.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.558, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.950.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1945.clone.1, %mul.1944.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) + %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.870.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.950.clone.1, %div.871.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.64.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.870.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1105.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1284,14 +1284,14 @@ StackFrames %add.949.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.851.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.428.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.872.clone.1, %add.949.clone.1), metadata={op_name="multiply.59"} %div.869.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.951.clone.1, %multiply.428.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2018.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1369, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.948.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.869.clone.1, %mul.2018.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2017.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.2019.clone.1, %add.948.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1369, %mul.2017.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.189 = f32[128,4]{0,1:T(4,128)} multiply(%add.947.clone.1, %add.947.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1200 = f32[]{:T(128)} constant(0) - %reduce.217 = f32[]{:T(128)} reduce(%square.189, %constant.1200), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1942.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1370, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.948.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.869.clone.1, %mul.1942.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1941.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1943.clone.1, %add.948.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1370, %mul.1941.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.265 = f32[128,4]{0,1:T(4,128)} multiply(%add.947.clone.1, %add.947.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1202 = f32[]{:T(128)} constant(0) + %reduce.217 = f32[]{:T(128)} reduce(%square.265, %constant.1202), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1202), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.159 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.217, %add.947.clone.1, %add.950.clone.1, %add.951.clone.1, %reduce.221.clone.1) } @@ -1307,39 +1307,39 @@ StackFrames ROOT %reduce_sum.361 = f32[]{:T(128)} add(%reduce_sum.359, %reduce_sum.360), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.412 (param_0.1372: f32[128,4], param_1.1554: f32[], param_2.1316: f32[], param_3.923: f32[], param_4.559: f32[128,4], param_5.472: f32[], param_6.361: f32[4,128], param_7.201: pred[], param_8.118: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1372 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.412 (param_0.1373: f32[128,4], param_1.1561: f32[], param_2.1319: f32[], param_3.923: f32[], param_4.561: f32[128,4], param_5.473: f32[], param_6.363: f32[4,128], param_7.206: pred[], param_8.123: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1373 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.923 = f32[]{:T(128)S(6)} parameter(3) - %mul.2046.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.923), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.201 = pred[]{:T(512)S(6)} parameter(7) - %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.201), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.361 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.472 = f32[]{:T(128)} parameter(5) - %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1970.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.923), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.206 = pred[]{:T(512)S(6)} parameter(7) + %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.363 = f32[4,128]{1,0:T(4,128)} parameter(6) + %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.473 = f32[]{:T(128)} parameter(5) + %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.899.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.474.clone.1, %div.900.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.287.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.288.clone.1, %bitcast.474.clone.1, %div.899.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1122.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.866.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1122.clone.1), dimensions={}, metadata={op_name="broadcast.78"} - %mul.2050.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.118 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1974.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.123 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1126.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.865.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1126.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.2049.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.118, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.968.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2050.clone.1, %mul.2049.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) - %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1973.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.123, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.968.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1974.clone.1, %mul.1973.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) + %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.70.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %select_n.287.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1125.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.864.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1125.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.2048.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.864.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.559 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1972.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.864.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.561 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1124.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.863.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1124.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.2047.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.559, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.967.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2048.clone.1, %mul.2047.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1554 = f32[]{:T(128)S(6)} parameter(1) - %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1554), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1971.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.967.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1972.clone.1, %mul.1971.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) + %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.894.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.967.clone.1, %div.895.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.67.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.894.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1123.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1347,35 +1347,35 @@ StackFrames %add.966.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.67.clone.1, %broadcast.861.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.431.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.896.clone.1, %add.966.clone.1), metadata={op_name="multiply.56"} %div.893.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.968.clone.1, %multiply.431.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2045.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1372, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.965.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.893.clone.1, %mul.2045.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2044.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.2046.clone.1, %add.965.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1372, %mul.2044.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.190 = f32[128,4]{0,1:T(4,128)} multiply(%add.964.clone.1, %add.964.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1203 = f32[]{:T(128)} constant(0) - %reduce.218 = f32[]{:T(128)} reduce(%square.190, %constant.1203), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1203), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.1969.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1373, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.965.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.893.clone.1, %mul.1969.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1968.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1970.clone.1, %add.965.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1373, %mul.1968.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.266 = f32[128,4]{0,1:T(4,128)} multiply(%add.964.clone.1, %add.964.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1205 = f32[]{:T(128)} constant(0) + %reduce.218 = f32[]{:T(128)} reduce(%square.266, %constant.1205), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.160 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.218, %add.964.clone.1, %add.967.clone.1, %add.968.clone.1, %reduce.222.clone.1) } -%fused_computation.421 (param_0.1200: f32[4,128], param_1.1320: f32[4,128]) -> f32[4,128] { - %param_0.1200 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1320 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.421 (param_0.1201: f32[4,128], param_1.1323: f32[4,128]) -> f32[4,128] { + %param_0.1201 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1323 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.1045 = f32[]{:T(128)} constant(0.00048828125) %broadcast.837 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1045), dimensions={}, metadata={op_name="broadcast.399"} - %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1320, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1323, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1043 = f32[]{:T(128)} constant(1e-06) %add.935 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1043), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.934 = f32[4,128]{1,0:T(4,128)} add(%div.767, %add.935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %rsqrt.168 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.934), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} %div.754 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.168, %add.934), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1040 = f32[]{:T(128)} constant(-0.5) - %mul.1995 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1040), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1986 = f32[4,128]{1,0:T(4,128)} multiply(%div.754, %mul.1995), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1985 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1200, %mul.1986), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1919 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1040), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1910 = f32[4,128]{1,0:T(4,128)} multiply(%div.754, %mul.1919), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1909 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1201, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.1039 = f32[]{:T(128)} constant(0.0009765625) - %mul.1994 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1039), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - ROOT %mul.1984 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1985, %mul.1994), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1918 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1039), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + ROOT %mul.1908 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1909, %mul.1918), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} } %region_0.1 (reduce_sum.137: s32[], reduce_sum.138: s32[]) -> s32[] { @@ -1384,31 +1384,31 @@ StackFrames ROOT %reduce_sum.139 = s32[]{:T(128)} add(%reduce_sum.137, %reduce_sum.138), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.425 (param_0.1219: pred[4,128]) -> s32[] { - %param_0.1219 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1219), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} +%fused_computation.425 (param_0.1220: pred[4,128]) -> s32[] { + %param_0.1220 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1220), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} %constant.1066 = s32[]{:T(128)} constant(0) ROOT %reduce.220 = s32[]{:T(128)} reduce(%convert_element_type.1403, %constant.1066), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.428 (param_0.1202: f32[4,128]) -> f32[4,128] { - %param_0.1202 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.428 (param_0.1203: f32[4,128]) -> f32[4,128] { + %param_0.1203 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1046 = f32[]{:T(128)} constant(0.00048828125) %broadcast.829 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1046), dimensions={}, metadata={op_name="broadcast.399"} - %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1202, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1203, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1044 = f32[]{:T(128)} constant(1e-06) %add.924 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1044), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.921 = f32[4,128]{1,0:T(4,128)} add(%div.759, %add.924), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} ROOT %rsqrt.166 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.921), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} } -%fused_computation.429 (param_0.1203: pred[4,128], param_1.1568: f32[]) -> f32[4,128] { - %param_0.1203 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} - %constant.1221 = f32[]{:T(128)} constant(0) - %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1221), dimensions={}, metadata={op_name="broadcast.99"} - ROOT %mul.1996 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1203, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.429 (param_0.1204: pred[4,128], param_1.1575: f32[]) -> f32[4,128] { + %param_0.1204 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %param_1.1575 = f32[]{:T(128)S(6)} parameter(1) + %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1575), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} + %constant.1223 = f32[]{:T(128)} constant(0) + %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.99"} + ROOT %mul.1920 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1204, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} } %fused_computation.431 () -> f32[64] { @@ -1417,43 +1417,43 @@ StackFrames %iota.46 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/layers/iota" stack_frame_id=0} %constant.1048 = s32[]{:T(128)} constant(2) %broadcast.839 = s32[64]{0:T(128)} broadcast(%constant.1048), dimensions={}, metadata={op_name="broadcast.391"} - %mul.1997 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.839), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} - %convert_element_type.1404 = f32[64]{0:T(128)} convert(%mul.1997), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %mul.1921 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.839), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} + %convert_element_type.1404 = f32[64]{0:T(128)} convert(%mul.1921), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %constant.1047 = f32[]{:T(128)} constant(0.0078125) %broadcast.838 = f32[64]{0:T(128)} broadcast(%constant.1047), dimensions={}, metadata={op_name="broadcast.392"} %div.768 = f32[64]{0:T(128)} multiply(%convert_element_type.1404, %broadcast.838), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.840, %div.768), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} } -%fused_computation.432 (param_0.1217: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.1217 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1217), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} +%fused_computation.432 (param_0.1218: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { + %param_0.1218 = s32[4,128]{1,0:T(4,128)} parameter(0) + %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1218), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %bitcast.418 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} ROOT %tuple.162 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.418, %convert_element_type.1405) } -%fused_computation.435 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { - %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.533 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.533) +%fused_computation.435 (param_0.1360: f32[2048,4]) -> bf16[4,2048] { + %param_0.1360 = f32[2048,4]{0,1:T(4,128)} parameter(0) + %bitcast.531 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.531) } -%fused_computation.436 (param_0.1358: f32[2048,4]) -> bf16[4,2048] { - %param_0.1358 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.532 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.532) +%fused_computation.436 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { + %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) + %bitcast.530 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.530) } -%fused_computation.437 (param_0.1360: f32[128,4]) -> bf16[4,128] { - %param_0.1360 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.534 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.534) +%fused_computation.437 (param_0.1361: f32[128,4]) -> bf16[4,128] { + %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) + %bitcast.532 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.532) } -%fused_computation.438 (param_0.1361: f32[128,4]) -> bf16[4,128] { - %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.535 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.535) +%fused_computation.438 (param_0.1362: f32[128,4]) -> bf16[4,128] { + %param_0.1362 = f32[128,4]{0,1:T(4,128)} parameter(0) + %bitcast.533 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1362), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.533) } %region_8.11 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { @@ -1462,40 +1462,40 @@ StackFrames ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.287.clone.clone (param_0.1345: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1345 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1345), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%fused_computation.287.clone.clone (param_0.1346: bf16[151936,2048]) -> bf16[151936,2048,1] { + %param_0.1346 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.526 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1346), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} } -%fused_computation.368.clone.clone (param_0.1346: f32[4,128], param_1.1535: bf16[4,128,2048], param_2.1278: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1278 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1278), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1535 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1432 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1535), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1346 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2135 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1346), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2134 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1432, %mul.2135), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1431 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2134), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.474 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.475, %convert_element_type.1431), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.368.clone.clone (param_0.1347: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1281: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1281 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.476 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1281), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1438 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1347 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2067 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1347), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.2066 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1438, %mul.2067), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2066), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.476, %convert_element_type.1437), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.439 (param_0.1362: bf16[151936,2048], param_1.1544: f32[4,128], param_2.1302: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { - %param_1.1544 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1302 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) +%fused_computation.439 (param_0.1363: bf16[151936,2048], param_1.1551: f32[4,128], param_2.1305: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { + %param_1.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1305 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) %param_3.913 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1544, %param_2.1302, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1362 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %fusion.252.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1362), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.269.clone.1, %fusion.252.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %constant.1193 = bf16[]{:T(256)} constant(-inf) - %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1193), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %fusion.270.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1551, %param_2.1305, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1363 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %fusion.253.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1363), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.270.clone.1, %fusion.253.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %constant.1195 = bf16[]{:T(256)} constant(-inf) + %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1195), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} ROOT %tuple.164 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,151936]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.223, %convolution.85.clone.1) } -%fused_computation.440 (param_0.1357: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.1357 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %bitcast.531 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1357), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.531) +%fused_computation.440 (param_0.1358: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { + %param_0.1358 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %bitcast.529 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.529) } %convert_element_type.767.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { @@ -1504,13 +1504,13 @@ StackFrames ROOT %add.755 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.155.clone.clone (param_0.1533: bf16[4,2048], param_1.1680: s32[]) -> bf16[2048] { - %param_0.1533 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1680 = s32[]{:T(128)S(6)} parameter(1) - %constant.1359 = s32[]{:T(128)} constant(0) - %dynamic_slice.384 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1533, %param_1.1680, %constant.1359), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1360 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.384, %constant.1360), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.155.clone.clone (param_0.1534: bf16[4,2048], param_1.1687: s32[]) -> bf16[2048] { + %param_0.1534 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) + %constant.1361 = s32[]{:T(128)} constant(0) + %dynamic_slice.388 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1687, %constant.1361), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1362 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.388, %constant.1362), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %region_14.16 (reduce_sum.204: f32[], reduce_sum.205: f32[]) -> f32[] { @@ -1519,25 +1519,25 @@ StackFrames ROOT %reduce_sum.206 = f32[]{:T(128)} add(%reduce_sum.204, %reduce_sum.205), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.58.clone.clone (param_0.1534: bf16[4,4,128,2048], param_1.1681: s32[]) -> f32[4,128] { - %param_0.1534 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1681 = s32[]{:T(128)S(6)} parameter(1) - %constant.1361 = s32[]{:T(128)} constant(0) - %dynamic_slice.385 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1681, %constant.1361, %constant.1361, %constant.1361), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.635 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1558 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.635), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.204 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1558, %convert_element_type.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1362 = f32[]{:T(128)} constant(0) - ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.204, %constant.1362), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} -} - -%fused_computation.179.clone.1.clone (param_0.1535: f32[4,128]) -> f32[4,128] { - %param_0.1535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1364 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1364), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1535, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1363 = f32[]{:T(128)} constant(1e-06) - %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1363), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} +%fused_computation.58.clone.clone (param_0.1535: bf16[4,4,128,2048], param_1.1688: s32[]) -> f32[4,128] { + %param_0.1535 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1688 = s32[]{:T(128)S(6)} parameter(1) + %constant.1363 = s32[]{:T(128)} constant(0) + %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1535, %param_1.1688, %constant.1363, %constant.1363, %constant.1363), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1564 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.633), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.280 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1564, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1364 = f32[]{:T(128)} constant(0) + ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.280, %constant.1364), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} +} + +%fused_computation.179.clone.1.clone (param_0.1536: f32[4,128]) -> f32[4,128] { + %param_0.1536 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %constant.1366 = f32[]{:T(128)} constant(0.00048828125) + %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1366), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1536, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1365 = f32[]{:T(128)} constant(1e-06) + %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1365), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.1039 = f32[4,128]{1,0:T(4,128)} add(%div.999, %closed_call.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.181 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1039), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } @@ -1548,158 +1548,158 @@ StackFrames ROOT %reduce_sum.212 = f32[]{:T(128)} add(%reduce_sum.207, %reduce_sum.211), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1549: bf16[4,2048,16,128], param_1.1691: s32[]) -> bf16[2048,16,128,1] { - %param_0.1549 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1691 = s32[]{:T(128)S(6)} parameter(1) - %constant.1375 = s32[]{:T(128)} constant(0) - %dynamic_slice.391 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1691, %constant.1375, %constant.1375, %constant.1375), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.646 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.114.clone.clone.clone.clone (param_0.1550: f32[4,128], param_1.1692: bf16[4,4,128,2048], param_2.1403: s32[], param_3.983: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.983 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.983), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1692 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1403 = s32[]{:T(128)S(6)} parameter(2) - %constant.1376 = s32[]{:T(128)} constant(0) - %dynamic_slice.392 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1692, %param_2.1403, %constant.1376, %constant.1376, %constant.1376), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.648 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1569 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.648), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1550 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2324 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2323 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1569, %mul.2324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1568 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.569 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.570, %convert_element_type.1568), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.647 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.61.clone.clone (param_0.1551: bf16[4,2048,16,128], param_1.1693: s32[], param_2.1404: f32[4,128], param_3.984: bf16[4,4,128,2048], param_4.603: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { - %param_2.1404 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.984 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1693 = s32[]{:T(128)S(6)} parameter(1) - %param_4.603 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1404, %param_3.984, %param_1.1693, %param_4.603), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1551 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1551, %param_1.1693), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.46.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1570 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.46.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.206 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1570, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1377 = f32[]{:T(128)} constant(0) - %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.206, %constant.1377), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.46.clone.3) -} - -%fused_computation.151.clone.1.clone (param_0.1552: f32[4,128,16]) -> f32[4,128,16] { - %param_0.1552 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1378 = f32[]{:T(128)} constant(0.0078125) - %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1378), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1552, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1379 = f32[]{:T(128)} constant(1e-06) - %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1379), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} +%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1550: bf16[4,2048,16,128], param_1.1698: s32[]) -> bf16[2048,16,128,1] { + %param_0.1550 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1698 = s32[]{:T(128)S(6)} parameter(1) + %constant.1377 = s32[]{:T(128)} constant(0) + %dynamic_slice.395 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1550, %param_1.1698, %constant.1377, %constant.1377, %constant.1377), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.644 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.114.clone.clone.clone.clone (param_0.1551: f32[4,128], param_1.1699: bf16[4,4,128,2048], param_2.1405: s32[], param_3.982: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.982 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.571 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1699 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1405 = s32[]{:T(128)S(6)} parameter(2) + %constant.1378 = s32[]{:T(128)} constant(0) + %dynamic_slice.396 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1699, %param_2.1405, %constant.1378, %constant.1378, %constant.1378), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.646 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.646), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2256 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1551), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2255 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %mul.2256), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2255), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.571, %convert_element_type.1574), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.645 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.61.clone.clone (param_0.1552: bf16[4,2048,16,128], param_1.1700: s32[], param_2.1406: f32[4,128], param_3.983: bf16[4,4,128,2048], param_4.604: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { + %param_2.1406 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.983 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1700 = s32[]{:T(128)S(6)} parameter(1) + %param_4.604 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1406, %param_3.983, %param_1.1700, %param_4.604), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1552 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1552, %param_1.1700), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.44.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %convert_element_type.1576 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.44.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.282 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1576, %convert_element_type.1576), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1379 = f32[]{:T(128)} constant(0) + %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.282, %constant.1379), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.44.clone.3) +} + +%fused_computation.151.clone.1.clone (param_0.1553: f32[4,128,16]) -> f32[4,128,16] { + %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) + %constant.1380 = f32[]{:T(128)} constant(0.0078125) + %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1380), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1553, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1381 = f32[]{:T(128)} constant(1e-06) + %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} %add.1043 = f32[4,128,16]{1,2,0:T(8,128)} add(%div.1001, %add.1044), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.183 = f32[4,128,16]{1,2,0:T(8,128)S(1)} rsqrt(%add.1043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.182.clone.clone (param_0.1548: bf16[4,128], param_1.1690: s32[]) -> bf16[128] { - %param_0.1548 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1690 = s32[]{:T(128)S(6)} parameter(1) - %constant.1374 = s32[]{:T(128)} constant(0) - %dynamic_slice.390 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1548, %param_1.1690, %constant.1374), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.645 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.182.clone.clone (param_0.1549: bf16[4,128], param_1.1697: s32[]) -> bf16[128] { + %param_0.1549 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) + %constant.1376 = s32[]{:T(128)} constant(0) + %dynamic_slice.394 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1697, %constant.1376), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.643 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.394), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.121.clone.1.clone (param_0.1553: f32[4,128,16], param_1.1694: bf16[4,128,16,128], param_2.1405: bf16[128]) -> bf16[4,128,16,128] { - %param_2.1405 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1405), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1694 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1572 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1694), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2326 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1553), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2325 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1572, %mul.2326), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1571 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.571 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.572, %convert_element_type.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.121.clone.1.clone (param_0.1554: f32[4,128,16], param_1.1701: bf16[4,128,16,128], param_2.1407: bf16[128]) -> bf16[4,128,16,128] { + %param_2.1407 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1407), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1701 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1578 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1554 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) + %mul.2258 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1554), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2257 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1578, %mul.2258), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1577 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2257), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.573, %convert_element_type.1577), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.90.clone.clone (param_0.1554: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { - %param_0.1554 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1554), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.90.clone.clone (param_0.1555: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { + %param_0.1555 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.129 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1554), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.209 = (bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) } %fused_computation.187.clone.clone () -> f32[64] { - %constant.1353 = f32[]{:T(128)} constant(1e+06) - %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1353), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %constant.1355 = f32[]{:T(128)} constant(1e+06) + %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1355), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} - %constant.1352 = s32[]{:T(128)} constant(2) - %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1352), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2310 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1556 = f32[64]{0:T(128)} convert(%mul.2310), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %constant.1354 = f32[]{:T(128)} constant(0.0078125) - %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1556, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1354 = s32[]{:T(128)} constant(2) + %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %mul.2242 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1562 = f32[64]{0:T(128)} convert(%mul.2242), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %constant.1356 = f32[]{:T(128)} constant(0.0078125) + %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1356), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1562, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.104, %div.995), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} } -%fused_computation.143.clone.clone (param_0.1528: f32[64], param_1.1676: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1676 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1676), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1528 = f32[64]{0:T(128)S(1)} parameter(0) - %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1528), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} +%fused_computation.143.clone.clone (param_0.1529: f32[64], param_1.1683: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1683 = f32[4,128]{1,0:T(4,128)} parameter(1) + %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1683), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %param_0.1529 = f32[64]{0:T(128)S(1)} parameter(0) + %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1529), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %div.996 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.998, %div.997), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1557 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1563 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} %convert_element_type.1189.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1557, %convert_element_type.1189.clone.3) + ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1563, %convert_element_type.1189.clone.3) } -%fused_computation.146.clone.1.clone (param_0.1529: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1529 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1355 = bf16[]{:T(256)} constant(-inf) - %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1529, %constant.1355), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1529, %constant.1355), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} +%fused_computation.146.clone.1.clone (param_0.1530: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1530 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) + %constant.1357 = bf16[]{:T(256)} constant(-inf) + %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.53 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.69, %pad.68), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.632 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.630 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.145.clone.1.clone (param_0.1544: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1544 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1372 = bf16[]{:T(256)} constant(-inf) - %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1544, %constant.1372), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1544, %constant.1372), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} +%fused_computation.145.clone.1.clone (param_0.1545: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1545 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) + %constant.1374 = bf16[]{:T(256)} constant(-inf) + %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.54 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.71, %pad.70), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.643 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.641 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.94.clone.clone (param_0.1555: bf16[4,128,16,64], param_1.1695: bf16[4,128,16,64], param_2.1406: bf16[4,128,128], param_3.985: bf16[4,128,128], param_4.604: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.383: bf16[128]) -> bf16[4,16,128,128] { - %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.94.clone.clone (param_0.1556: bf16[4,128,16,64], param_1.1702: bf16[4,128,16,64], param_2.1408: bf16[4,128,128], param_3.984: bf16[4,128,128], param_4.605: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.384: bf16[128]) -> bf16[4,16,128,128] { + %param_6.384 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.575 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.384), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} %param_5.499 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1574 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.604 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2333 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.604), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2332 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1574, %mul.2333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2332), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %convert_element_type.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.985 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2331 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.985), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2329 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.573, %mul.2331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1695 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1380 = bf16[]{:T(256)} constant(-inf) - %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1695, %constant.1380), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1555 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1555, %constant.1380), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %convert_element_type.1580 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_4.605 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) + %mul.2265 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.605), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2264 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1580, %mul.2265), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1579 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.575, %convert_element_type.1579), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.984 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2263 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.984), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2261 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %mul.2263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1702 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1382 = bf16[]{:T(256)} constant(-inf) + %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1702, %constant.1382), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1556 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1556, %constant.1382), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.56 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.75, %pad.74), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1406 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2330 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1406), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2328 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2329, %mul.2328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %constant.1381 = bf16[]{:T(256)} constant(0.08838) - %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2327 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - ROOT %bitcast.649 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_2.1408 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2262 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1408), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2260 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2262), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2261, %mul.2260), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + %constant.1383 = bf16[]{:T(256)} constant(0.08838) + %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1383), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %mul.2259 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.647 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2259), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } %region_16.18 (reduce_sum.213: f32[], reduce_sum.214: f32[]) -> f32[] { @@ -1708,159 +1708,159 @@ StackFrames ROOT %reduce_sum.218 = f32[]{:T(128)} add(%reduce_sum.213, %reduce_sum.214), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1540: bf16[4,2048,8,128], param_1.1685: s32[]) -> bf16[2048,8,128,1] { - %param_0.1540 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) - %constant.1367 = s32[]{:T(128)} constant(0) - %dynamic_slice.388 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1540, %param_1.1685, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.640 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.113.clone.clone.clone.clone (param_0.1541: f32[4,128], param_1.1686: bf16[4,4,128,2048], param_2.1399: s32[], param_3.980: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.980 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.980), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1686 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) - %constant.1368 = s32[]{:T(128)} constant(0) - %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1686, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.642 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1562 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.642), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1541 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2314 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1541), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2313 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1562, %mul.2314), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1561 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2313), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.564, %convert_element_type.1561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.641 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.84.clone.clone (param_0.1542: bf16[4,2048,8,128], param_1.1687: s32[], param_2.1400: f32[4,128], param_3.981: bf16[4,4,128,2048], param_4.601: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { - %param_2.1400 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.981 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) - %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1400, %param_3.981, %param_1.1687, %param_4.601), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1542 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1542, %param_1.1687), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1541: bf16[4,2048,8,128], param_1.1692: s32[]) -> bf16[2048,8,128,1] { + %param_0.1541 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1692 = s32[]{:T(128)S(6)} parameter(1) + %constant.1369 = s32[]{:T(128)} constant(0) + %dynamic_slice.392 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1541, %param_1.1692, %constant.1369, %constant.1369, %constant.1369), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.638 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.113.clone.clone.clone.clone (param_0.1542: f32[4,128], param_1.1693: bf16[4,4,128,2048], param_2.1401: s32[], param_3.979: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.979 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.979), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1693 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1401 = s32[]{:T(128)S(6)} parameter(2) + %constant.1370 = s32[]{:T(128)} constant(0) + %dynamic_slice.393 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1693, %param_2.1401, %constant.1370, %constant.1370, %constant.1370), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.640 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1568 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.640), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1542 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2246 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1542), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2245 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1568, %mul.2246), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.565, %convert_element_type.1567), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.639 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.84.clone.clone (param_0.1543: bf16[4,2048,8,128], param_1.1694: s32[], param_2.1402: f32[4,128], param_3.980: bf16[4,4,128,2048], param_4.602: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { + %param_2.1402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.980 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1694 = s32[]{:T(128)S(6)} parameter(1) + %param_4.602 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1402, %param_3.980, %param_1.1694, %param_4.602), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1543 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1543, %param_1.1694), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %convolution.50.clone.3 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.73.clone.3, %fusion.87.clone.3), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1563 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.205 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1563, %convert_element_type.1563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1369 = f32[]{:T(128)} constant(0) - %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.205, %constant.1369), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %convert_element_type.1569 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.281 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1569, %convert_element_type.1569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1371 = f32[]{:T(128)} constant(0) + %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.281, %constant.1371), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.206 = (f32[4,128,8]{1,2,0:T(8,128)S(1)}, bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.246, %convolution.50.clone.3) } -%fused_computation.154.clone.1.clone (param_0.1543: f32[4,128,8]) -> f32[4,128,8] { - %param_0.1543 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1370 = f32[]{:T(128)} constant(0.0078125) - %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1370), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1543, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1371 = f32[]{:T(128)} constant(1e-06) - %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1371), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} +%fused_computation.154.clone.1.clone (param_0.1544: f32[4,128,8]) -> f32[4,128,8] { + %param_0.1544 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) + %constant.1372 = f32[]{:T(128)} constant(0.0078125) + %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1372), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1544, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1373 = f32[]{:T(128)} constant(1e-06) + %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1373), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} %add.1040 = f32[4,128,8]{1,2,0:T(8,128)} add(%div.1000, %add.1041), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.182 = f32[4,128,8]{1,2,0:T(8,128)S(1)} rsqrt(%add.1040), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.184.clone.clone (param_0.1527: bf16[4,128], param_1.1675: s32[]) -> bf16[128] { - %param_0.1527 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1675 = s32[]{:T(128)S(6)} parameter(1) - %constant.1351 = s32[]{:T(128)} constant(0) - %dynamic_slice.381 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1527, %param_1.1675, %constant.1351), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.631 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.139.clone.1.clone (param_0.1545: f32[4,128,8], param_1.1688: bf16[4,128,8,128], param_2.1401: bf16[128]) -> bf16[4,128,8,128] { - %param_2.1401 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1401), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1688 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1565 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1688), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1545 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2316 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1545), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2315 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1565, %mul.2316), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1564 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2315), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.565 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.566, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.126.clone.clone (param_0.1546: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1546 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1546), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.184.clone.clone (param_0.1528: bf16[4,128], param_1.1682: s32[]) -> bf16[128] { + %param_0.1528 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) + %constant.1353 = s32[]{:T(128)} constant(0) + %dynamic_slice.385 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1528, %param_1.1682, %constant.1353), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.629 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.139.clone.1.clone (param_0.1546: f32[4,128,8], param_1.1695: bf16[4,128,8,128], param_2.1403: bf16[128]) -> bf16[4,128,8,128] { + %param_2.1403 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1403), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1695 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1571 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1695), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1546 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) + %mul.2248 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1546), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2247 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1571, %mul.2248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1570 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2247), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.567, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.126.clone.clone (param_0.1547: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { + %param_0.1547 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1546), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.207 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) } -%fused_computation.129.clone.clone (param_0.1547: bf16[4,128,8,64], param_1.1689: bf16[4,128,8,64], param_2.1402: bf16[4,128,128], param_3.982: bf16[4,128,128], param_4.602: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.382: bf16[128]) -> bf16[4,8,128,128] { - %param_6.382 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.382), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.129.clone.clone (param_0.1548: bf16[4,128,8,64], param_1.1696: bf16[4,128,8,64], param_2.1404: bf16[4,128,128], param_3.981: bf16[4,128,128], param_4.603: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.383: bf16[128]) -> bf16[4,8,128,128] { + %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.569 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} %param_5.498 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1567 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.602 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2322 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.602), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2321 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1567, %mul.2322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %convert_element_type.1566), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.982 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2320 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2318 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.567, %mul.2320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1689 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1373 = bf16[]{:T(256)} constant(-inf) - %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1689, %constant.1373), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1547 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1547, %constant.1373), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %convert_element_type.1573 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_4.603 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) + %mul.2254 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.603), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2253 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1573, %mul.2254), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1572 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2253), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.569, %convert_element_type.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.981 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2252 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.981), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2250 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %mul.2252), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1696 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1375 = bf16[]{:T(256)} constant(-inf) + %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1696, %constant.1375), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1548 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1548, %constant.1375), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.55 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.73, %pad.72), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1402 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2319 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1402), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2317 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2319), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2318, %mul.2317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.644 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_2.1404 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2251 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1404), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2249 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2250, %mul.2249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.642 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.169.clone.clone (param_0.1536: bf16[4,2048,8,128], param_1.1682: s32[]) -> bf16[1,2048,8,128] { - %param_0.1536 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) - %constant.1365 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.386 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1536, %param_1.1682, %constant.1365, %constant.1365, %constant.1365), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} -} - -%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1537: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { - %param_0.1537 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.208 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1537), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.636 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.208), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.111.clone.clone.clone.clone (param_0.1538: f32[4,128], param_1.1683: bf16[4,4,128,2048], param_2.1397: s32[], param_3.978: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.978 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.978), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1683 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1397 = s32[]{:T(128)S(6)} parameter(2) - %constant.1366 = s32[]{:T(128)} constant(0) - %dynamic_slice.387 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1683, %param_2.1397, %constant.1366, %constant.1366, %constant.1366), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.638 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1560 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.638), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1538 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2312 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1538), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2311 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1560, %mul.2312), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1559 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2311), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.561 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.562, %convert_element_type.1559), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.637 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.140.clone.clone (param_0.1539: bf16[1,2048,8,128], param_1.1684: f32[4,128], param_2.1398: bf16[4,4,128,2048], param_3.979: s32[], param_4.600: bf16[2048]) -> bf16[4,8,128,128] { - %param_1.1684 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1398 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.979 = s32[]{:T(128)S(6)} parameter(3) - %param_4.600 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.372 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1684, %param_2.1398, %param_3.979, %param_4.600), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1539 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.371 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1539), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.372, %fusion.371), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.639 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.188.clone.clone (param_0.1577: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { - %param_0.1577 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1577), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.662 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.662, %slice.11) +%fused_computation.169.clone.clone (param_0.1537: bf16[4,2048,8,128], param_1.1689: s32[]) -> bf16[1,2048,8,128] { + %param_0.1537 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) + %param_1.1689 = s32[]{:T(128)S(6)} parameter(1) + %constant.1367 = s32[]{:T(128)} constant(0) + ROOT %dynamic_slice.390 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1537, %param_1.1689, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1538: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { + %param_0.1538 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %copy.204 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1538), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} + ROOT %bitcast.634 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.204), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.111.clone.clone.clone.clone (param_0.1539: f32[4,128], param_1.1690: bf16[4,4,128,2048], param_2.1399: s32[], param_3.977: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.977 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.977), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1690 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) + %constant.1368 = s32[]{:T(128)} constant(0) + %dynamic_slice.391 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1690, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.636 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1566 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.636), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1539 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2244 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2243 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1566, %mul.2244), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2243), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.563, %convert_element_type.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.635 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.562), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.140.clone.clone (param_0.1540: bf16[1,2048,8,128], param_1.1691: f32[4,128], param_2.1400: bf16[4,4,128,2048], param_3.978: s32[], param_4.601: bf16[2048]) -> bf16[4,8,128,128] { + %param_1.1691 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1400 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.978 = s32[]{:T(128)S(6)} parameter(3) + %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.373 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1691, %param_2.1400, %param_3.978, %param_4.601), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1540 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.372 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1540), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.373, %fusion.372), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + ROOT %bitcast.637 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.188.clone.clone (param_0.1578: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { + %param_0.1578 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) + %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1578), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.660 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} + ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.660, %slice.11) } %region_17.20 (reduce_sum.219: f32[], reduce_sum.220: f32[]) -> f32[] { @@ -1869,36 +1869,36 @@ StackFrames ROOT %reduce_sum.221 = f32[]{:T(128)} add(%reduce_sum.219, %reduce_sum.220), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1556: bf16[4,16,128,2048], param_1.1696: s32[]) -> bf16[16,128,2048,1] { - %param_0.1556 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1696 = s32[]{:T(128)S(6)} parameter(1) - %constant.1382 = s32[]{:T(128)} constant(0) - %dynamic_slice.393 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1556, %param_1.1696, %constant.1382, %constant.1382, %constant.1382), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.650 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,2048], param_1.1703: s32[]) -> bf16[16,128,2048,1] { + %param_0.1557 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) + %constant.1384 = s32[]{:T(128)} constant(0) + %dynamic_slice.397 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1557, %param_1.1703, %constant.1384, %constant.1384, %constant.1384), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.648 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.397), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,128]) -> bf16[4,128,16,128] { - %param_0.1557 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.651 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1558: bf16[4,16,128,128]) -> bf16[4,128,16,128] { + %param_0.1558 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.649 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.64.clone.clone (param_0.1558: bf16[4,16,128,2048], param_1.1697: s32[], param_2.1407: bf16[4,16,128,128], param_3.986: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { - %param_3.986 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.64.clone.clone (param_0.1559: bf16[4,16,128,2048], param_1.1704: s32[], param_2.1409: bf16[4,16,128,128], param_3.985: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { + %param_3.985 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1704 = s32[]{:T(128)S(6)} parameter(1) %constant.436.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.240.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.986, %param_1.1697, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.240.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1407 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1407), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1558 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1558, %param_1.1697), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.242.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.985, %param_1.1704, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.242.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %param_2.1409 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1409), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_0.1559 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1559, %param_1.1704), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %convolution.62.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.96.clone.3, %fusion.95.clone.3), window={size=1x16}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %bitcast.203.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %add.768.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.227.clone.3, %bitcast.203.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.207 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %convert_element_type.1575), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1383 = f32[]{:T(128)} constant(0) - %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.207, %constant.1383), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %convert_element_type.1581 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.283 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1581, %convert_element_type.1581), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1385 = f32[]{:T(128)} constant(0) + %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.283, %constant.1385), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.210 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.248, %add.768.clone.3) } @@ -1908,93 +1908,93 @@ StackFrames ROOT %add.754 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.156.clone.clone (param_0.1530: bf16[4,2048], param_1.1677: s32[]) -> bf16[2048] { - %param_0.1530 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1677 = s32[]{:T(128)S(6)} parameter(1) - %constant.1356 = s32[]{:T(128)} constant(0) - %dynamic_slice.382 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1530, %param_1.1677, %constant.1356), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1357 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.382, %constant.1357), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.156.clone.clone (param_0.1531: bf16[4,2048], param_1.1684: s32[]) -> bf16[2048] { + %param_0.1531 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1684 = s32[]{:T(128)S(6)} parameter(1) + %constant.1358 = s32[]{:T(128)} constant(0) + %dynamic_slice.386 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1684, %constant.1358), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1359 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.386, %constant.1359), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.13.clone.clone.clone (param_0.1531: bf16[4,6144,2048], param_1.1678: s32[]) -> bf16[6144,2048,1] { - %param_0.1531 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1678 = s32[]{:T(128)S(6)} parameter(1) - %constant.1358 = s32[]{:T(128)} constant(0) - %dynamic_slice.383 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1678, %constant.1358, %constant.1358), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.634 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.13.clone.clone.clone (param_0.1532: bf16[4,6144,2048], param_1.1685: s32[]) -> bf16[6144,2048,1] { + %param_0.1532 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) + %constant.1360 = s32[]{:T(128)} constant(0) + %dynamic_slice.387 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1532, %param_1.1685, %constant.1360, %constant.1360), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.632 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %bitcast_fusion.1.clone.clone (bitcast_input.4: bf16[4,128,2048]) -> bf16[4,128,2048] { - %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) -} - -%fused_computation.14.clone.clone (param_0.1532: bf16[4,128,2048], param_1.1679: bf16[4,6144,2048], param_2.1396: s32[]) -> bf16[6144,4,128] { - %param_1.1679 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1396 = s32[]{:T(128)S(6)} parameter(2) - %fusion.369 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1679, %param_2.1396), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1532 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.370 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1532), kind=kLoop, calls=%bitcast_fusion.1.clone.clone - ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.369, %fusion.370), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.180.clone.1.clone (param_0.1559: f32[4,128]) -> f32[4,128] { - %param_0.1559 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1385 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1385), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1559, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1384 = f32[]{:T(128)} constant(1e-06) - %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1384), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.631 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) +} + +%fused_computation.14.clone.clone (param_0.1533: bf16[4,128,2048], param_1.1686: bf16[4,6144,2048], param_2.1398: s32[]) -> bf16[6144,4,128] { + %param_1.1686 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1398 = s32[]{:T(128)S(6)} parameter(2) + %fusion.370 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1686, %param_2.1398), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1533 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.371 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1533), kind=kLoop, calls=%bitcast_fusion.1.clone.clone + ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.370, %fusion.371), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.180.clone.1.clone (param_0.1560: f32[4,128]) -> f32[4,128] { + %param_0.1560 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %constant.1387 = f32[]{:T(128)} constant(0.00048828125) + %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1387), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1560, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1386 = f32[]{:T(128)} constant(1e-06) + %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1386), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.1046 = f32[4,128]{1,0:T(4,128)} add(%div.1002, %closed_call.110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.184 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1046), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.12.clone.1.clone.clone (param_0.1563: bf16[4,2048,6144], param_1.1701: s32[]) -> bf16[2048,6144,1] { - %param_0.1563 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1701 = s32[]{:T(128)S(6)} parameter(1) - %constant.1387 = s32[]{:T(128)} constant(0) - %dynamic_slice.395 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1563, %param_1.1701, %constant.1387, %constant.1387), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.119.clone.3.clone.clone (param_0.1564: f32[4,128], param_1.1702: bf16[4,128,2048], param_2.1410: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1410 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1410), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1702 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1579 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1564 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2337 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1564), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2336 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1579, %mul.2337), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2336), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.577 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.578, %convert_element_type.1578), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.21.clone.clone (param_0.1565: bf16[4,2048,6144], param_1.1703: s32[], param_2.1411: f32[4,128], param_3.988: bf16[4,128,2048], param_4.606: bf16[2048]) -> bf16[4,128,6144] { - %param_2.1411 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.988 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.606 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.376 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1411, %param_3.988, %param_4.606), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1565 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) - %fusion.375 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1565, %param_1.1703), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.376, %fusion.375), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.11.clone.1.clone.clone (param_0.1567: bf16[4,2048,6144], param_1.1705: s32[]) -> bf16[2048,6144,1] { - %param_0.1567 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1705 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.12.clone.1.clone.clone (param_0.1564: bf16[4,2048,6144], param_1.1708: s32[]) -> bf16[2048,6144,1] { + %param_0.1564 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1708 = s32[]{:T(128)S(6)} parameter(1) %constant.1389 = s32[]{:T(128)} constant(0) - %dynamic_slice.396 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1567, %param_1.1705, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.655 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.47.clone.1.clone.clone (param_0.1566: bf16[6144,4,128], param_1.1704: bf16[4,128,6144]) -> bf16[4,128,6144] { - %param_1.1704 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1388 = bf16[]{:T(256)} constant(1) - %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1388), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %dynamic_slice.399 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1564, %param_1.1708, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.651 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.119.clone.3.clone.clone (param_0.1565: f32[4,128], param_1.1709: bf16[4,128,2048], param_2.1412: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1412 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.579 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1412), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1709 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1585 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1565 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2269 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1565), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2268 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1585, %mul.2269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1584 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.579, %convert_element_type.1584), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.21.clone.clone (param_0.1566: bf16[4,2048,6144], param_1.1710: s32[], param_2.1413: f32[4,128], param_3.987: bf16[4,128,2048], param_4.607: bf16[2048]) -> bf16[4,128,6144] { + %param_2.1413 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.987 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.607 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.377 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1413, %param_3.987, %param_4.607), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1566 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1710 = s32[]{:T(128)S(6)} parameter(1) + %fusion.376 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1566, %param_1.1710), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.377, %fusion.376), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.11.clone.1.clone.clone (param_0.1568: bf16[4,2048,6144], param_1.1712: s32[]) -> bf16[2048,6144,1] { + %param_0.1568 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1712 = s32[]{:T(128)S(6)} parameter(1) + %constant.1391 = s32[]{:T(128)} constant(0) + %dynamic_slice.400 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1568, %param_1.1712, %constant.1391, %constant.1391), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.400), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.47.clone.1.clone.clone (param_0.1567: bf16[6144,4,128], param_1.1711: bf16[4,128,6144]) -> bf16[4,128,6144] { + %param_1.1711 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1390 = bf16[]{:T(256)} constant(1) + %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1390), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} + %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1711), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.69 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.1047 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.1003 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.1047), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.2339 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1704, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %mul.2271 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1711, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} From 24b6b9bfb0d3d026508602d87ff76225dabcddbc Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 6 Jun 2026 18:23:29 +0000 Subject: [PATCH 12/52] ci: fix GPU unit and integration test failures on multi-GPU runners --- .github/workflows/run_tests_against_package.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 9df299ab21..c2d36baf27 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -91,6 +91,8 @@ jobs: || (inputs.device_type == 'cpu' && 'tpu' || inputs.device_type) }} ALLOW_MULTIPLE_LIBTPU_LOAD: ${{ inputs.device_type == 'cpu' && 'true' || '' }} # bypass /tmp/libtpu_lockfile check for cpu tests, which don't actually use accelerators (to allow concurrency) + NCCL_P2P_DISABLE: ${{ inputs.device_type == 'cuda12' && '1' || '' }} + NCCL_IB_DISABLE: ${{ inputs.device_type == 'cuda12' && '1' || '' }} options: ${{ inputs.container_resource_option }} steps: - name: Checkout MaxText @@ -172,6 +174,11 @@ jobs: done fi fi + # Restrict GPU unit tests to a single GPU to avoid multi-device NCCL failures + if [[ "${INPUTS_PYTEST_MARKER}" == *"not integration_test"* ]]; then + export CUDA_VISIBLE_DEVICES=0 + echo "Restricting GPU unit tests to a single GPU: CUDA_VISIBLE_DEVICES=0" + fi fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist From df85a1039f36ca6e2905373c450bd3d709c11ede Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Mon, 8 Jun 2026 17:15:00 +0000 Subject: [PATCH 13/52] Revert GPU unit and integration test changes in run_tests_against_package.yml --- .github/workflows/run_tests_against_package.yml | 7 ------- 1 file changed, 7 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index c2d36baf27..9df299ab21 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -91,8 +91,6 @@ jobs: || (inputs.device_type == 'cpu' && 'tpu' || inputs.device_type) }} ALLOW_MULTIPLE_LIBTPU_LOAD: ${{ inputs.device_type == 'cpu' && 'true' || '' }} # bypass /tmp/libtpu_lockfile check for cpu tests, which don't actually use accelerators (to allow concurrency) - NCCL_P2P_DISABLE: ${{ inputs.device_type == 'cuda12' && '1' || '' }} - NCCL_IB_DISABLE: ${{ inputs.device_type == 'cuda12' && '1' || '' }} options: ${{ inputs.container_resource_option }} steps: - name: Checkout MaxText @@ -174,11 +172,6 @@ jobs: done fi fi - # Restrict GPU unit tests to a single GPU to avoid multi-device NCCL failures - if [[ "${INPUTS_PYTEST_MARKER}" == *"not integration_test"* ]]; then - export CUDA_VISIBLE_DEVICES=0 - echo "Restricting GPU unit tests to a single GPU: CUDA_VISIBLE_DEVICES=0" - fi fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist From 439aec6abc1c51ce96b5579917937e7c56026f6f Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Mon, 8 Jun 2026 18:29:37 +0000 Subject: [PATCH 14/52] chore: keep aqtp dependency to debug GPU/NCCL failure --- src/dependencies/requirements/base_requirements/requirements.txt | 1 + .../requirements/generated_requirements/cuda12-requirements.txt | 1 + .../generated_requirements/tpu-post-train-requirements.txt | 1 + .../requirements/generated_requirements/tpu-requirements.txt | 1 + src/dependencies/requirements/requirements.txt | 1 + .../requirements/requirements_decoupled_jax_0_7.1.txt | 1 + 6 files changed, 6 insertions(+) diff --git a/src/dependencies/requirements/base_requirements/requirements.txt b/src/dependencies/requirements/base_requirements/requirements.txt index 919f2665f7..5ba8ee5093 100644 --- a/src/dependencies/requirements/base_requirements/requirements.txt +++ b/src/dependencies/requirements/base_requirements/requirements.txt @@ -1,4 +1,5 @@ absl-py +aqtp array-record chex cloud-accelerator-diagnostics diff --git a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt index 96bd058585..4e52d9125a 100644 --- a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt @@ -10,6 +10,7 @@ annotated-doc>=0.0.4 annotated-types>=0.7.0 antlr4-python3-runtime>=4.9.3 anyio>=4.13.0 +aqtp>=0.9.0 array-record>=0.8.3 astroid>=4.0.4 astunparse>=1.6.3 diff --git a/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt index 2bae5fd262..b0bd183f81 100644 --- a/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/tpu-post-train-requirements.txt @@ -14,6 +14,7 @@ antlr4-python3-runtime>=4.9.3 anyio>=4.13.0 apache-tvm-ffi>=0.1.11 appnope>=0.1.4 ; sys_platform == 'darwin' +aqtp>=0.9.0 array-record>=0.8.3 astor>=0.8.1 astroid>=4.0.4 diff --git a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt index 6e0474b7ff..26ba4fcdda 100644 --- a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt @@ -10,6 +10,7 @@ annotated-doc>=0.0.4 annotated-types>=0.7.0 antlr4-python3-runtime>=4.9.3 anyio>=4.13.0 +aqtp>=0.9.0 array-record>=0.8.3 astroid>=4.0.4 astunparse>=1.6.3 diff --git a/src/dependencies/requirements/requirements.txt b/src/dependencies/requirements/requirements.txt index 633dfec057..05c2be074b 100644 --- a/src/dependencies/requirements/requirements.txt +++ b/src/dependencies/requirements/requirements.txt @@ -1,4 +1,5 @@ absl-py +aqtp array-record cloud-accelerator-diagnostics cloud-tpu-diagnostics diff --git a/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt b/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt index e1cec8bef7..8f904a3641 100644 --- a/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt +++ b/src/dependencies/requirements/requirements_decoupled_jax_0_7.1.txt @@ -1,4 +1,5 @@ absl_py>=2.3.1 +aqtp>=0.9.0 chex>=0.1.90 datasets>=4.2.0 etils>=1.13.0 From 65c788d37f8fcd8a9a826df347eb49de072fb137 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Mon, 8 Jun 2026 20:26:06 +0000 Subject: [PATCH 15/52] ci: disable NCCL P2P and IB for stable GPU tests --- .github/workflows/run_tests_against_package.yml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 9df299ab21..a13f777619 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -172,6 +172,13 @@ jobs: done fi fi + + # Set NCCL environment variables to prevent multi-GPU communication failures on the runner. + # Disabling P2P and IB forces NCCL to fall back to shared memory and standard socket communication, + # which is stable and avoids "invalid argument" errors on virtualized runner environments. + export NCCL_P2P_DISABLE=1 + export NCCL_IB_DISABLE=1 + echo "Setting NCCL_P2P_DISABLE=1 and NCCL_IB_DISABLE=1 for GPU execution." fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist From eca0a49b941a6339c7d8a78731b9a6fe0f95b45c Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Mon, 8 Jun 2026 21:41:02 +0000 Subject: [PATCH 16/52] ci: improve GPU CUDA library discovery and add diagnostics --- .../workflows/run_tests_against_package.yml | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index a13f777619..ec258b6eca 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -162,16 +162,22 @@ jobs: # Dynamically discover the 'nvidia' folder and prepend all its sub-library # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent # JAX from partially loading incompatible system-level CUDA libraries. - if [ -d ".venv/lib" ]; then - NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1) - if [ -n "${NVIDIA_DIR}" ]; then - for dir in "${NVIDIA_DIR}"/*; do - if [ -d "$dir/lib" ]; then - export LD_LIBRARY_PATH=$(pwd)/$dir/lib:${LD_LIBRARY_PATH} - fi - done - fi + NVIDIA_DIR=$(find -L .venv -name "nvidia" -type d 2>/dev/null | head -n 1) + echo "=== GPU Diagnostics ===" + echo "Found NVIDIA_DIR: ${NVIDIA_DIR}" + if [ -n "${NVIDIA_DIR}" ]; then + for dir in "${NVIDIA_DIR}"/*; do + if [ -d "$dir/lib" ]; then + ABS_LIB_PATH=$(realpath "$dir/lib") + export LD_LIBRARY_PATH=${ABS_LIB_PATH}:${LD_LIBRARY_PATH} + echo "Prepended to LD_LIBRARY_PATH: ${ABS_LIB_PATH}" + fi + done + else + echo "WARNING: nvidia directory not found under .venv!" fi + echo "Final LD_LIBRARY_PATH: ${LD_LIBRARY_PATH}" + echo "=======================" # Set NCCL environment variables to prevent multi-GPU communication failures on the runner. # Disabling P2P and IB forces NCCL to fall back to shared memory and standard socket communication, From feb8715dc2f52b59c48a5646f20dde972cf66125 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 16:05:16 +0000 Subject: [PATCH 17/52] ci: fix GPU CUDA library discovery to use site-packages/nvidia --- .github/workflows/run_tests_against_package.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index ec258b6eca..b2e5c9e0f9 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -162,7 +162,7 @@ jobs: # Dynamically discover the 'nvidia' folder and prepend all its sub-library # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent # JAX from partially loading incompatible system-level CUDA libraries. - NVIDIA_DIR=$(find -L .venv -name "nvidia" -type d 2>/dev/null | head -n 1) + NVIDIA_DIR=$(find -L .venv -path "*/site-packages/nvidia" -type d 2>/dev/null | head -n 1) echo "=== GPU Diagnostics ===" echo "Found NVIDIA_DIR: ${NVIDIA_DIR}" if [ -n "${NVIDIA_DIR}" ]; then From 43bf7306de3e32687e6be1211fc1b52ea10cb11c Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 16:27:09 +0000 Subject: [PATCH 18/52] ci: re-enable NCCL P2P and IB for native GPU communication --- .github/workflows/run_tests_against_package.yml | 7 ------- 1 file changed, 7 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index b2e5c9e0f9..0d78c934f6 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -178,13 +178,6 @@ jobs: fi echo "Final LD_LIBRARY_PATH: ${LD_LIBRARY_PATH}" echo "=======================" - - # Set NCCL environment variables to prevent multi-GPU communication failures on the runner. - # Disabling P2P and IB forces NCCL to fall back to shared memory and standard socket communication, - # which is stable and avoids "invalid argument" errors on virtualized runner environments. - export NCCL_P2P_DISABLE=1 - export NCCL_IB_DISABLE=1 - echo "Setting NCCL_P2P_DISABLE=1 and NCCL_IB_DISABLE=1 for GPU execution." fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist From 74145d94fea9678eb4b5bcda8832f96c57ecc41c Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 16:54:20 +0000 Subject: [PATCH 19/52] ci: configure stable NCCL vars for containerized multi-GPU runs --- .github/workflows/run_tests_against_package.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 0d78c934f6..abc76932db 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -178,6 +178,12 @@ jobs: fi echo "Final LD_LIBRARY_PATH: ${LD_LIBRARY_PATH}" echo "=======================" + + # Configure NCCL for stable single-node communication in Docker containers + export NCCL_SOCKET_IFNAME=lo + export NCCL_NET_GDR_LEVEL=0 + export NCCL_DEBUG=INFO + echo "Set NCCL_SOCKET_IFNAME=lo, NCCL_NET_GDR_LEVEL=0, and NCCL_DEBUG=INFO for GPU execution." fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist From d4dacafdb00c522645a25a9e99bd4d98de797830 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 17:40:16 +0000 Subject: [PATCH 20/52] chore: update use_qwix_quantization to false and update error message --- src/maxtext/configs/base.yml | 2 +- src/maxtext/configs/types.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index bb5619b236..f4c7acd3af 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -138,7 +138,7 @@ save_quantized_params_path: "" # when left as is, corresponds to training # accepted values are "inference" model_call_mode: "" -use_qwix_quantization: true # [DEPRECATED: AQT will be removed in a future release. It is strongly recommended to set use_qwix_quantization to true] whether to use qwix for quantization. if set to true, the model will be quantized using qwix. +use_qwix_quantization: false # [DEPRECATED: AQT will be removed in a future release. It is strongly recommended to set use_qwix_quantization to true] whether to use qwix for quantization. if set to true, the model will be quantized using qwix. use_manual_quantization: false # a flag if to use manual quantization for batch split. Only used if use_batch_split_schedule is true. # quantization calibration method used for weights and activations. supported methods can be found in https://github.com/google/qwix/blob/dc2a0770351c740e5ab3cce7c0efe9f7beacce9e/qwix/qconfig.py#l70-l80 weight_quantization_calibration_method: "absmax" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 60b60a67df..13b5089994 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2646,7 +2646,7 @@ def get_num_target_devices(): raise ValueError( f"Quantization type '{self.quantization}' without Qwix (use_qwix_quantization=False) " f"is unsupported because legacy AQT has been completely removed. " - f"Please migrate to Qwix by setting use_qwix_quantization=True." + f"Please migrate to Qwix by setting use_qwix_quantization=False." ) # Default quantization sharding count to number of local devices if not set. From 4271648c38f04eda39c5cf03b0f8ea7d3df75ce9 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 17:46:08 +0000 Subject: [PATCH 21/52] Making changes to types and reverting use_qwix_quantization to False --- src/maxtext/configs/types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 13b5089994..60f0eddd0f 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -430,7 +430,7 @@ class Quantization(BaseModel): kv_quant_axis: KvQuantAxis = Field(KvQuantAxis.HEADS_AND_DKV, description="Axes to quantize over for the KV cache.") kv_quant_dtype: Literal["int8", "int4"] = Field("int8", description="Data type for KV cache quantization.") quantization_local_shard_count: int = Field(-1, description="Shards the range finding operation for quantization.") - use_qwix_quantization: bool = Field(True, description="Whether to use qwix for quantization.") + use_qwix_quantization: bool = Field(False, description="Whether to use qwix for quantization.") use_manual_quantization: bool = Field( False, description="Whether to use manual quantization for batch split. Only used if use_batch_split_schedule is True.", From 7d24a3e7a99ee8d56e53d04db62ae73ee36a57e4 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 21:19:11 +0000 Subject: [PATCH 22/52] Fix quantization test failures by defaulting use_qwix_quantization to True - Change default of use_qwix_quantization to True in types.py and base.yml. - Fix typo in validator error message. - Update test_quantization_fallbacks to use fp8_gpu and use_qwix_quantization=False to safely test fallback path. TAG=agy CONV=9aac9cad-26d1-453f-9b81-c70a14dd59dc --- src/maxtext/configs/base.yml | 2 +- src/maxtext/configs/types.py | 4 ++-- tests/unit/quantizations_test.py | 3 ++- 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index f4c7acd3af..f8f4f8ae22 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -138,7 +138,7 @@ save_quantized_params_path: "" # when left as is, corresponds to training # accepted values are "inference" model_call_mode: "" -use_qwix_quantization: false # [DEPRECATED: AQT will be removed in a future release. It is strongly recommended to set use_qwix_quantization to true] whether to use qwix for quantization. if set to true, the model will be quantized using qwix. +use_qwix_quantization: true # whether to use qwix for quantization. if set to true, the model will be quantized using qwix. use_manual_quantization: false # a flag if to use manual quantization for batch split. Only used if use_batch_split_schedule is true. # quantization calibration method used for weights and activations. supported methods can be found in https://github.com/google/qwix/blob/dc2a0770351c740e5ab3cce7c0efe9f7beacce9e/qwix/qconfig.py#l70-l80 weight_quantization_calibration_method: "absmax" diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 60f0eddd0f..60b60a67df 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -430,7 +430,7 @@ class Quantization(BaseModel): kv_quant_axis: KvQuantAxis = Field(KvQuantAxis.HEADS_AND_DKV, description="Axes to quantize over for the KV cache.") kv_quant_dtype: Literal["int8", "int4"] = Field("int8", description="Data type for KV cache quantization.") quantization_local_shard_count: int = Field(-1, description="Shards the range finding operation for quantization.") - use_qwix_quantization: bool = Field(False, description="Whether to use qwix for quantization.") + use_qwix_quantization: bool = Field(True, description="Whether to use qwix for quantization.") use_manual_quantization: bool = Field( False, description="Whether to use manual quantization for batch split. Only used if use_batch_split_schedule is True.", @@ -2646,7 +2646,7 @@ def get_num_target_devices(): raise ValueError( f"Quantization type '{self.quantization}' without Qwix (use_qwix_quantization=False) " f"is unsupported because legacy AQT has been completely removed. " - f"Please migrate to Qwix by setting use_qwix_quantization=False." + f"Please migrate to Qwix by setting use_qwix_quantization=True." ) # Default quantization sharding count to number of local devices if not set. diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index a8270bc2de..037c9e6f2c 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -433,7 +433,8 @@ def test_quantization_fallbacks(self): # Cover the fallback return None path in _get_quant_config when an unsupported scheme is passed config_invalid = pyconfig.initialize( [None, get_test_config_path()], - quantization="int4", + quantization="fp8_gpu", + use_qwix_quantization=False, ) self.assertIsNone(quantizations.configure_quantization(config_invalid)) From c86e9286c4498daeedca50b040353a34e4aabcaa Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 21:29:00 +0000 Subject: [PATCH 23/52] Fix formatting in maxengine.py to satisfy pyink --- src/maxtext/inference/maxengine/maxengine.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index ed0b5ff74a..1703395819 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -536,7 +536,6 @@ def unapply_adapter(self, base_params, adapter_config, adapter_params): else: lora_utils.unapply_lora_from_base_params(base_params, adapter_params, lora_scale_factor) - def _maybe_stack_prefill_result_cache(self, cache): """Stack the caches across the layers.""" if not self.config.stack_prefill_result_cache: From dba254ac1c95ae2d3d3378762e7070711e78642a Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 22:24:35 +0000 Subject: [PATCH 24/52] Delete obsolete quantization tests following AQT deprecation --- tests/unit/aqt_serve_roundtrip_nnx_test.py | 170 ------------------ tests/unit/layerwise_quantization_nnx_test.py | 77 -------- 2 files changed, 247 deletions(-) delete mode 100644 tests/unit/aqt_serve_roundtrip_nnx_test.py delete mode 100644 tests/unit/layerwise_quantization_nnx_test.py diff --git a/tests/unit/aqt_serve_roundtrip_nnx_test.py b/tests/unit/aqt_serve_roundtrip_nnx_test.py deleted file mode 100644 index 00d6825de7..0000000000 --- a/tests/unit/aqt_serve_roundtrip_nnx_test.py +++ /dev/null @@ -1,170 +0,0 @@ -# Copyright 2023-2026 Google LLC -# -# 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 -# -# https://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. - -"""Round-trip test for the NNX serve-mode AQT checkpoint path. - -Builds a small NNX model in CONVERT mode with int8 quantization, runs a forward -to populate `qrhs.frozen`, saves the serve-mode-shape state to a local orbax -checkpoint, then reloads via `from_pretrained(quant_mode_str="serve")` and -checks that the loaded QTensor leaves match what was saved. - -This guards the chain of issues exercised by serve-mode reload (sharding helper -for QTensor, v[...] vs get_value() for composite values, Param-only filter -dropping aqt-typed leaves, Partitioned-unwrap for matching on-disk paths). -""" - -import os -import sys -import tempfile -import unittest - -import jax -import jax.numpy as jnp -import orbax.checkpoint as ocp -from flax import nnx -from flax.core.meta import Partitioned -from flax.linen import partitioning as nn_partitioning - -from maxtext.configs import pyconfig -from maxtext.utils import maxtext_utils, model_creation_utils, maxtext_utils_nnx -from maxtext.utils.globals import MAXTEXT_PKG_DIR -from maxtext.utils.layerwise_quantization import LayerwiseQuantization - - -def _wrap_value(node): - """Add `{"value": ...}` per-leaf wrap matching `_load_and_quantize_nnx` save format.""" - if isinstance(node, dict): - return {k: _wrap_value(v) for k, v in node.items()} - return {"value": node} - - -def _unbox(x): - return x.value if isinstance(x, Partitioned) else x - - -def _walk_qrhs(state): - """Yield (path_str, variable) pairs for every qrhs.frozen entry in an nnx.State.""" - for path, var in state.flat_state(): - keys = [str(getattr(k, "key", k)) for k in path] - if "qrhs" in keys and "frozen" in keys: - yield ".".join(keys), var - - -class ServeModeRoundTripTest(unittest.TestCase): - """End-to-end save+reload of a serve-mode NNX AQT checkpoint.""" - - def _init_cfg(self, ckpt_path, *, checkpoint_is_quantized): - """Build a pyconfig for save or reload.""" - # Use base.yml + gpt3-52k. The decoupled test config strips - # logical_axis_rules (e.g. "norm"), which the AQT serve-mode model - # construction needs. - base_yml = os.path.join(MAXTEXT_PKG_DIR, "configs", "base.yml") - args = [ - sys.argv[0], - base_yml, - "model_name=gpt3-52k", - "pure_nnx=true", - "enable_nnx=true", - "pure_nnx_decoder=true", - "max_target_length=64", - "max_prefill_predict_length=16", - "per_device_batch_size=1", - "scan_layers=true", - "quantization=int8", - "checkpoint_storage_use_ocdbt=false", - "checkpoint_storage_use_zarr3=false", - "skip_jax_distributed_system=true", - ] - if checkpoint_is_quantized: - args += [ - f"load_parameters_path={ckpt_path}", - "checkpoint_is_quantized=true", - "enable_checkpointing=true", # required by config validator when load_parameters_path is set - ] - else: - args += ["enable_checkpointing=false"] - return pyconfig.initialize(args) - - def test_save_then_reload_preserves_qrhs_frozen(self): - """Save a serve-mode-shape NNX checkpoint, then reload it and compare qvalue arrays.""" - with tempfile.TemporaryDirectory() as tmpdir: - ckpt_path = os.path.join(tmpdir, "quantized_ckpt") - - # Step 1: build CONVERT-mode model + run forward to populate qrhs.frozen. - cfg_save = self._init_cfg(ckpt_path, checkpoint_is_quantized=False) - mesh = maxtext_utils.get_mesh_from_config(cfg_save) - rngs = maxtext_utils_nnx.create_nnx_rngs(cfg_save) - with nn_partitioning.axis_rules(cfg_save.logical_axis_rules): - convert_model = model_creation_utils.from_config( - cfg_save, - mesh=mesh, - rngs=rngs, - model_mode="train", - quant_mode_str="convert", - ) - L = cfg_save.max_prefill_predict_length - tokens = jnp.zeros((1, L), dtype=jnp.int32) - pos = jnp.arange(L, dtype=jnp.int32)[None, :] - seg = jnp.ones((1, L), dtype=jnp.int32) - with nn_partitioning.axis_rules(cfg_save.logical_axis_rules): - _ = convert_model(tokens, pos, decoder_segment_ids=seg, enable_dropout=False, model_mode="train") - - # Step 2: capture the qrhs.frozen leaves we expect to round-trip, then save. - convert_state = nnx.state(convert_model).to_pure_dict() - serve_state = LayerwiseQuantization._strip_kernels_at_quantized_paths(convert_state) # pylint: disable=protected-access - saved_qrhs = {} - for path, var in _walk_qrhs(nnx.state(convert_model)): - qt = var.value if hasattr(var, "value") else var - saved_qrhs[path] = _unbox(qt.qvalue) - - # Replicate arrays across the mesh; orbax rejects SingleDeviceSharding - # once another test has initialized JAX-distributed state. - replicated = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec()) - serve_state = jax.tree.map( - lambda x: jax.device_put(x, replicated) if isinstance(x, jax.Array) else x, - serve_state, - ) - - orbax_checkpointer = ocp.PyTreeCheckpointer(use_ocdbt=False, use_zarr3=False) - orbax_checkpointer.save(ckpt_path, _wrap_value(serve_state), force=True) - self.assertGreater(len(saved_qrhs), 0, "Test config must produce at least one qrhs.frozen leaf") - - # Step 3: reload via from_pretrained in serve mode. - cfg_load = self._init_cfg(ckpt_path, checkpoint_is_quantized=True) - with nn_partitioning.axis_rules(cfg_load.logical_axis_rules): - loaded_model = model_creation_utils.from_pretrained( - cfg_load, - mesh=mesh, - model_mode="autoregressive", - quant_mode_str="serve", - ) - - # Step 4: assert every saved qrhs.frozen leaf matches what was persisted. - loaded_state = nnx.state(loaded_model) - loaded_qrhs = dict(_walk_qrhs(loaded_state)) - self.assertEqual(set(saved_qrhs.keys()), set(loaded_qrhs.keys())) - for path, saved_qv in saved_qrhs.items(): - var = loaded_qrhs[path] - qt = var.value if hasattr(var, "value") else var - loaded_qv = _unbox(qt.qvalue) - self.assertEqual(loaded_qv.shape, saved_qv.shape, f"shape mismatch at {path}") - self.assertEqual(loaded_qv.dtype, saved_qv.dtype, f"dtype mismatch at {path}") - self.assertTrue( - jnp.array_equal(loaded_qv.astype(jnp.int32), saved_qv.astype(jnp.int32)), - f"qvalue not preserved at {path}", - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/layerwise_quantization_nnx_test.py b/tests/unit/layerwise_quantization_nnx_test.py deleted file mode 100644 index bbd43f7964..0000000000 --- a/tests/unit/layerwise_quantization_nnx_test.py +++ /dev/null @@ -1,77 +0,0 @@ -# Copyright 2023–2026 Google LLC -# -# 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 -# -# https://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. - -"""Unit tests for the NNX path of layerwise_quantization. - -Covers `_strip_kernels_at_quantized_paths` — the convert→serve shape converter -that drops the redundant full-precision kernel from quantized DenseGeneral -nodes while leaving non-quantized kernels (norms, embeddings) intact. -""" - -import unittest - -from maxtext.utils.layerwise_quantization import LayerwiseQuantization - - -class StripKernelsTest(unittest.TestCase): - - def test_drops_kernel_at_quantized_dense(self): - """A node with both `kernel` and `AqtDotGeneral_0` loses the kernel.""" - state = { - "decoder": { - "layers": { - "mlp": { - "wi": { - "kernel": "FULL_PRECISION_W", - "AqtDotGeneral_0": {"qrhs": {"frozen": "AQT_STATE"}}, - } - } - } - } - } - out = LayerwiseQuantization._strip_kernels_at_quantized_paths(state) # pylint: disable=protected-access - wi = out["decoder"]["layers"]["mlp"]["wi"] - self.assertNotIn("kernel", wi) - self.assertIn("AqtDotGeneral_0", wi) - self.assertEqual(wi["AqtDotGeneral_0"]["qrhs"]["frozen"], "AQT_STATE") - - def test_preserves_non_quantized_kernel(self): - """A non-quantized kernel (e.g. embedding, norm) survives.""" - state = { - "decoder": { - "decoder_norm": {"scale": "NORM_SCALE"}, - "logits_dense": {"kernel": "LOGITS_KERNEL"}, # no AqtDotGeneral_0 sibling - }, - "token_embedder": {"embedding": "EMB"}, - } - out = LayerwiseQuantization._strip_kernels_at_quantized_paths(state) # pylint: disable=protected-access - self.assertEqual(out["decoder"]["logits_dense"]["kernel"], "LOGITS_KERNEL") - self.assertEqual(out["decoder"]["decoder_norm"]["scale"], "NORM_SCALE") - self.assertEqual(out["token_embedder"]["embedding"], "EMB") - - def test_mixed_tree(self): - """Quantized + non-quantized at the same depth: only the quantized one strips.""" - state = { - "self_attention": { - "qkv_proj": {"kernel": "QKV", "AqtDotGeneral_0": "AQT"}, - "out": {"kernel": "OUT_FULL"}, # non-quantized output projection - } - } - out = LayerwiseQuantization._strip_kernels_at_quantized_paths(state) # pylint: disable=protected-access - self.assertNotIn("kernel", out["self_attention"]["qkv_proj"]) - self.assertEqual(out["self_attention"]["out"]["kernel"], "OUT_FULL") - - -if __name__ == "__main__": - unittest.main() From 431210d06c537a59de535f4d2c3375b8954eb475 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 22:48:18 +0000 Subject: [PATCH 25/52] Fix NNX + quantization in MaxEngine: remove NotImplementedError and fix mode detection for Qwix --- src/maxtext/inference/maxengine/maxengine.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 1703395819..162865072b 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -122,7 +122,7 @@ def __init__(self, config: Any, devices: Any | None = None): # quant enabled we stay in `train` mode and let AQT quantize per-forward # against the full-precision kernel — same numerical result as `serve` # for absmax calibration, just slower. - nnx_quant_mode_str = "serve" if (quant is not None and config.checkpoint_is_quantized) else "train" + nnx_quant_mode_str = "serve" if (config.quantization and config.checkpoint_is_quantized) else "train" # We need both PREFILL and AR abstract models because the cache vars inherit # CACHE_BATCH_PREFILL vs CACHE_BATCH from the construction model_mode, and # bulk_insert searches for the substring "cache_batch" in the AR-mode names. @@ -384,8 +384,6 @@ def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwar def _load_params_nnx(self, params, rng): """NNX equivalent of load_params: returns an nnx.Param state and populates KV cache shardings.""" - if self.config.quantization: - raise NotImplementedError("pure_nnx + quantization not yet supported. Use pure_nnx=False.") if params: print("Resharding given NNX params") From ce9f91f38dc67d30482cdb33ef1f8f376fa4ca44 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Tue, 9 Jun 2026 22:50:58 +0000 Subject: [PATCH 26/52] Fix pylint error in attention_compressed.py: migrate AqtQuantization to Quantization --- src/maxtext/layers/attention_compressed.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/maxtext/layers/attention_compressed.py b/src/maxtext/layers/attention_compressed.py index e9a25f46b5..ac66b62981 100644 --- a/src/maxtext/layers/attention_compressed.py +++ b/src/maxtext/layers/attention_compressed.py @@ -38,7 +38,7 @@ from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned from maxtext.layers.linears import DenseGeneral, DeepSeekV4GroupedLinear from maxtext.layers.normalizations import RMSNorm -from maxtext.layers.quantizations import AqtQuantization as Quant +from maxtext.layers.quantizations import Quantization as Quant from maxtext.inference.kvcache import KVQuant From 93b41b4022eaacbb3b4331d0a9a10a93ea011fb6 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Wed, 10 Jun 2026 01:43:35 +0000 Subject: [PATCH 27/52] Update reference HLO for deepseek3 TAG=agy CONV=39f31ff3-8d14-4d1b-b519-217aeb07b904 --- tests/utils/reference_hlo_deepseek3.txt | 1948 +++++++++++------------ 1 file changed, 974 insertions(+), 974 deletions(-) diff --git a/tests/utils/reference_hlo_deepseek3.txt b/tests/utils/reference_hlo_deepseek3.txt index ffff9103ad..cc3471559c 100644 --- a/tests/utils/reference_hlo_deepseek3.txt +++ b/tests/utils/reference_hlo_deepseek3.txt @@ -10,21 +10,21 @@ StackFrames %region_46.56 (top_k.25: bf16[], top_k.26: bf16[], top_k.27: s32[], top_k.28: s32[]) -> pred[] { - %constant.1408 = s32[]{:T(128)} constant(0) - %constant.1409 = s32[]{:T(128)} constant(2147483647) + %constant.1424 = s32[]{:T(128)} constant(0) + %constant.1425 = s32[]{:T(128)} constant(2147483647) %top_k.25 = bf16[]{:T(256)} parameter(0), metadata={op_name="top_k"} %top_k.26 = bf16[]{:T(256)} parameter(1), metadata={op_name="top_k"} %top_k.27 = s32[]{:T(128)} parameter(2), metadata={op_name="top_k"} %top_k.28 = s32[]{:T(128)} parameter(3), metadata={op_name="top_k"} %convert.393 = f32[]{:T(128)S(6)} convert(%top_k.25), metadata={op_name="convert.18"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %bitcast-convert.39 = s32[]{:T(128)S(6)} bitcast-convert(%convert.393), metadata={op_name="bitcast-convert.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.144 = pred[]{:T(512)S(6)} compare(%bitcast-convert.39, %constant.1408), direction=LT, metadata={op_name="compare.38"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.40 = s32[]{:T(128)S(6)} xor(%constant.1409, %bitcast-convert.39), metadata={op_name="xor.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.144 = pred[]{:T(512)S(6)} compare(%bitcast-convert.39, %constant.1424), direction=LT, metadata={op_name="compare.38"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.40 = s32[]{:T(128)S(6)} xor(%constant.1425, %bitcast-convert.39), metadata={op_name="xor.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.127 = s32[]{:T(128)S(6)} select(%compare.144, %xor.40, %bitcast-convert.39), metadata={op_name="select.16"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %convert.394 = f32[]{:T(128)S(6)} convert(%top_k.26), metadata={op_name="convert.19"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %bitcast-convert.40 = s32[]{:T(128)S(6)} bitcast-convert(%convert.394), metadata={op_name="bitcast-convert.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.145 = pred[]{:T(512)S(6)} compare(%bitcast-convert.40, %constant.1408), direction=LT, metadata={op_name="compare.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.41 = s32[]{:T(128)S(6)} xor(%constant.1409, %bitcast-convert.40), metadata={op_name="xor.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.145 = pred[]{:T(512)S(6)} compare(%bitcast-convert.40, %constant.1424), direction=LT, metadata={op_name="compare.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.41 = s32[]{:T(128)S(6)} xor(%constant.1425, %bitcast-convert.40), metadata={op_name="xor.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.128 = s32[]{:T(128)S(6)} select(%compare.145, %xor.41, %bitcast-convert.40), metadata={op_name="select.17"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %compare.146 = pred[]{:T(512)S(6)} compare(%select.127, %select.128), direction=GT, metadata={op_name="compare.0"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %compare.147 = pred[]{:T(512)S(6)} compare(%select.128, %select.127), direction=GT, metadata={op_name="compare.117"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} @@ -78,19 +78,19 @@ StackFrames %region_107.126 (psum.6: bf16[], psum.9: bf16[]) -> bf16[] { %psum.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} %psum.9 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1407 = bf16[]{:T(256)} add(%psum.6, %psum.9), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.1417 = bf16[]{:T(256)} add(%psum.6, %psum.9), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %region_108.127 (psum.10: bf16[], psum.11: bf16[]) -> bf16[] { %psum.10 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} %psum.11 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1408 = bf16[]{:T(256)} add(%psum.10, %psum.11), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.1418 = bf16[]{:T(256)} add(%psum.10, %psum.11), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %region_109.128 (psum.14: bf16[], psum.15: bf16[]) -> bf16[] { %psum.14 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} %psum.15 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1409 = bf16[]{:T(256)} add(%psum.14, %psum.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.1419 = bf16[]{:T(256)} add(%psum.14, %psum.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %region_62.73 (reduce-window.111: s32[], reduce-window.112: s32[]) -> s32[] { @@ -211,167 +211,167 @@ StackFrames %param_0.17 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) %param_1.108 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.13 = s32[1024]{0:T(1024)} custom-call(%param_1.108), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %slice.920 = s32[512]{0:T(512)} slice(%custom-call.13), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %reshape.3318 = s32[4,128]{1,0:T(4,128)} reshape(%slice.920), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.847 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3318), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %gather.187 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} gather(%param_0.17, %transpose.847), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %transpose.846 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%gather.187), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - ROOT %reshape.3317 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.846), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %slice.892 = s32[512]{0:T(512)} slice(%custom-call.13), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %reshape.3298 = s32[4,128]{1,0:T(4,128)} reshape(%slice.892), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %transpose.847 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3298), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %gather.183 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} gather(%param_0.17, %transpose.847), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %transpose.846 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%gather.183), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + ROOT %reshape.3297 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.846), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } %fused_computation.6 (param_0.20: f32[163840,32], param_1.110: s32[1024]) -> f32[512,32] { - %param_0.20 = f32[163840,32]{1,0:T(8,128)} parameter(0) + %param_0.20 = f32[163840,32]{1,0:T(8,128)S(1)} parameter(0) %param_1.110 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.15 = s32[1024]{0:T(1024)} custom-call(%param_1.110), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %slice.922 = s32[512]{0:T(512)} slice(%custom-call.15), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %reshape.3326 = s32[4,128]{1,0:T(4,128)} reshape(%slice.922), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %transpose.853 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3326), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %gather.189 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.20, %transpose.853), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %transpose.852 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.189), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - ROOT %reshape.3325 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.852), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %slice.894 = s32[512]{0:T(512)} slice(%custom-call.15), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %reshape.3306 = s32[4,128]{1,0:T(4,128)} reshape(%slice.894), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} + %transpose.853 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3306), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} + %gather.185 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.20, %transpose.853), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %transpose.852 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.185), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + ROOT %reshape.3305 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.852), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} } %fused_computation.7 (param_0.23: f32[163840,32], param_1.112: s32[1024]) -> f32[512,32] { %param_0.23 = f32[163840,32]{1,0:T(8,128)} parameter(0) %param_1.112 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.17 = s32[1024]{0:T(1024)} custom-call(%param_1.112), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %slice.924 = s32[512]{0:T(512)} slice(%custom-call.17), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %reshape.3334 = s32[4,128]{1,0:T(4,128)} reshape(%slice.924), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %transpose.859 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3334), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %gather.191 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.23, %transpose.859), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %transpose.858 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.191), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - ROOT %reshape.3333 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.858), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %slice.896 = s32[512]{0:T(512)} slice(%custom-call.17), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %reshape.3314 = s32[4,128]{1,0:T(4,128)} reshape(%slice.896), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} + %transpose.859 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3314), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} + %gather.187 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.23, %transpose.859), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %transpose.858 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.187), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + ROOT %reshape.3313 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.858), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} } %fused_computation.8 (param_0.26: f32[163840,32], param_1.120: s32[1024]) -> f32[512,32] { %param_0.26 = f32[163840,32]{1,0:T(8,128)} parameter(0) %param_1.120 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.25 = s32[1024]{0:T(1024)} custom-call(%param_1.120), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %slice.932 = s32[512]{0:T(512)} slice(%custom-call.25), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %reshape.3342 = s32[4,128]{1,0:T(4,128)} reshape(%slice.932), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %transpose.865 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3342), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %gather.193 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.26, %transpose.865), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %transpose.864 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.193), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - ROOT %reshape.3341 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %slice.904 = s32[512]{0:T(512)} slice(%custom-call.25), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %reshape.3322 = s32[4,128]{1,0:T(4,128)} reshape(%slice.904), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} + %transpose.865 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3322), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} + %gather.189 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.26, %transpose.865), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %transpose.864 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.189), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + ROOT %reshape.3321 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} } %fused_computation.9 (param_0.29: f32[163840,32], param_1.122: s32[1024]) -> f32[512,32] { %param_0.29 = f32[163840,32]{1,0:T(8,128)} parameter(0) %param_1.122 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.27 = s32[1024]{0:T(1024)} custom-call(%param_1.122), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %slice.934 = s32[512]{0:T(512)} slice(%custom-call.27), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %reshape.3350 = s32[4,128]{1,0:T(4,128)} reshape(%slice.934), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %transpose.871 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3350), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %gather.195 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.29, %transpose.871), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %transpose.870 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.195), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - ROOT %reshape.3349 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.870), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %slice.906 = s32[512]{0:T(512)} slice(%custom-call.27), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %reshape.3330 = s32[4,128]{1,0:T(4,128)} reshape(%slice.906), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} + %transpose.871 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3330), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} + %gather.191 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.29, %transpose.871), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + %transpose.870 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.191), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} + ROOT %reshape.3329 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.870), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} } %fused_computation.10 (param_0.32: bf16[4096,512], param_1.126: s32[4096]) -> bf16[4096,512] { %param_0.32 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.126 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.31 = s32[4096]{0:T(1024)} custom-call(%param_1.126), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.938 = s32[4096]{0:T(1024)} slice(%custom-call.31), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3358 = s32[4096]{0:T(1024)} reshape(%slice.938), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.877 = s32[4096]{0:T(1024)} transpose(%reshape.3358), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.197 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.32, %transpose.877), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.876 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.197), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3357 = bf16[4096,512]{1,0:T(8,128)(2,1)} reshape(%transpose.876), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.910 = s32[4096]{0:T(1024)} slice(%custom-call.31), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3338 = s32[4096]{0:T(1024)} reshape(%slice.910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.877 = s32[4096]{0:T(1024)} transpose(%reshape.3338), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.193 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.32, %transpose.877), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.876 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.193), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3337 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.876), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.11 (param_0.35: bf16[4096,512], param_1.128: s32[4096]) -> bf16[4096,512] { %param_0.35 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.128 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.33 = s32[4096]{0:T(1024)} custom-call(%param_1.128), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.940 = s32[4096]{0:T(1024)} slice(%custom-call.33), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3366 = s32[4096]{0:T(1024)} reshape(%slice.940), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.883 = s32[4096]{0:T(1024)} transpose(%reshape.3366), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.199 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.35, %transpose.883), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.882 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.199), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3365 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.882), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.912 = s32[4096]{0:T(1024)} slice(%custom-call.33), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3346 = s32[4096]{0:T(1024)} reshape(%slice.912), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.883 = s32[4096]{0:T(1024)} transpose(%reshape.3346), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.195 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.35, %transpose.883), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.882 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.195), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3345 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.882), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.12 (param_0.38: bf16[4096,512], param_1.130: s32[4096]) -> bf16[4096,512] { %param_0.38 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.130 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.35 = s32[4096]{0:T(1024)} custom-call(%param_1.130), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.942 = s32[4096]{0:T(1024)} slice(%custom-call.35), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3374 = s32[4096]{0:T(1024)} reshape(%slice.942), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.889 = s32[4096]{0:T(1024)} transpose(%reshape.3374), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.201 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.38, %transpose.889), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.888 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.201), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3373 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.888), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.914 = s32[4096]{0:T(1024)} slice(%custom-call.35), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3354 = s32[4096]{0:T(1024)} reshape(%slice.914), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.889 = s32[4096]{0:T(1024)} transpose(%reshape.3354), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.197 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.38, %transpose.889), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.888 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.197), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3353 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.888), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.13 (param_0.41: bf16[4096,512], param_1.132: s32[4096]) -> bf16[4096,512] { %param_0.41 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.132 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.37 = s32[4096]{0:T(1024)} custom-call(%param_1.132), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.944 = s32[4096]{0:T(1024)} slice(%custom-call.37), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3382 = s32[4096]{0:T(1024)} reshape(%slice.944), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.895 = s32[4096]{0:T(1024)} transpose(%reshape.3382), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.203 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.41, %transpose.895), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.894 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.203), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3381 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.894), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.916 = s32[4096]{0:T(1024)} slice(%custom-call.37), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3362 = s32[4096]{0:T(1024)} reshape(%slice.916), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.895 = s32[4096]{0:T(1024)} transpose(%reshape.3362), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.199 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.41, %transpose.895), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.894 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.199), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3361 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.894), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %fused_computation.15 (param_0.47: s32[256], param_1.124: s32[1024]) -> s32[263] { %param_0.47 = s32[256]{0:T(256)S(1)} parameter(0) %param_1.124 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.29 = s32[1024]{0:T(1024)} custom-call(%param_1.124), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %slice.936 = s32[263]{0:T(512)} slice(%custom-call.29), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3413 = s32[263]{0:T(512)} reshape(%slice.936), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.911 = s32[263]{0:T(512)} transpose(%reshape.3413), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.208 = s32[263]{0:T(512)} gather(%param_0.47, %transpose.911), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.910 = s32[263]{0:T(512)} transpose(%gather.208), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3412 = s32[263]{0:T(512)S(1)} reshape(%transpose.910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %slice.908 = s32[263]{0:T(512)} slice(%custom-call.29), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.3393 = s32[263]{0:T(512)} reshape(%slice.908), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.911 = s32[263]{0:T(512)} transpose(%reshape.3393), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.204 = s32[263]{0:T(512)} gather(%param_0.47, %transpose.911), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.910 = s32[263]{0:T(512)} transpose(%gather.204), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.3392 = s32[263]{0:T(512)S(1)} reshape(%transpose.910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} } %fused_computation.16 (param_0.50: s32[256], param_1.134: s32[1024]) -> s32[263] { %param_0.50 = s32[256]{0:T(256)S(1)} parameter(0) %param_1.134 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.39 = s32[1024]{0:T(1024)} custom-call(%param_1.134), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %slice.946 = s32[263]{0:T(512)} slice(%custom-call.39), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3436 = s32[263]{0:T(512)} reshape(%slice.946), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.921 = s32[263]{0:T(512)} transpose(%reshape.3436), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.211 = s32[263]{0:T(512)} gather(%param_0.50, %transpose.921), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.920 = s32[263]{0:T(512)} transpose(%gather.211), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3435 = s32[263]{0:T(512)S(1)} reshape(%transpose.920), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %slice.918 = s32[263]{0:T(512)} slice(%custom-call.39), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.3416 = s32[263]{0:T(512)} reshape(%slice.918), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.921 = s32[263]{0:T(512)} transpose(%reshape.3416), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.207 = s32[263]{0:T(512)} gather(%param_0.50, %transpose.921), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.920 = s32[263]{0:T(512)} transpose(%gather.207), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.3415 = s32[263]{0:T(512)S(1)} reshape(%transpose.920), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} } %region_173.198.clone (scatter-add.94: bf16[], scatter-add.96: bf16[]) -> bf16[] { %scatter-add.94 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} %scatter-add.96 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - ROOT %add.1875 = bf16[]{:T(256)} add(%scatter-add.94, %scatter-add.96), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.1885 = bf16[]{:T(256)} add(%scatter-add.94, %scatter-add.96), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %fused_computation.21 (param_0.55: bf16[129280,512], param_1.65: s32[512], param_2.24: bf16[512,512]) -> bf16[129280,512] { %param_0.55 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) %param_1.65 = s32[512]{0:T(512)S(1)} parameter(1) - %reshape.3490 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.65), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.954 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3490), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %reshape.3470 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.65), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %transpose.954 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3470), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} %param_2.24 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} parameter(2) - %reshape.3491 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} reshape(%param_2.24), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} - %transpose.955 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%reshape.3491), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} - ROOT %scatter.77 = bf16[129280,512]{1,0:T(8,128)(2,1)} scatter(%param_0.55, %transpose.954, %transpose.955), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_173.198.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} + %reshape.3471 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} reshape(%param_2.24), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} + %transpose.955 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%reshape.3471), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} + ROOT %scatter.73 = bf16[129280,512]{1,0:T(8,128)(2,1)} scatter(%param_0.55, %transpose.954, %transpose.955), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_173.198.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} } %region_12.18 (top_k.0: bf16[], top_k.6: bf16[], top_k.7: s32[], top_k.8: s32[]) -> pred[] { - %constant.1369 = s32[]{:T(128)} constant(0) - %constant.1370 = s32[]{:T(128)} constant(2147483647) + %constant.1385 = s32[]{:T(128)} constant(0) + %constant.1386 = s32[]{:T(128)} constant(2147483647) %top_k.0 = bf16[]{:T(256)} parameter(0), metadata={op_name="top_k"} %top_k.6 = bf16[]{:T(256)} parameter(1), metadata={op_name="top_k"} %top_k.7 = s32[]{:T(128)} parameter(2), metadata={op_name="top_k"} %top_k.8 = s32[]{:T(128)} parameter(3), metadata={op_name="top_k"} %convert.385 = f32[]{:T(128)S(6)} convert(%top_k.0), metadata={op_name="convert.16"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %bitcast-convert.35 = s32[]{:T(128)S(6)} bitcast-convert(%convert.385), metadata={op_name="bitcast-convert.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.128 = pred[]{:T(512)S(6)} compare(%bitcast-convert.35, %constant.1369), direction=LT, metadata={op_name="compare.35"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.36 = s32[]{:T(128)S(6)} xor(%constant.1370, %bitcast-convert.35), metadata={op_name="xor.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.128 = pred[]{:T(512)S(6)} compare(%bitcast-convert.35, %constant.1385), direction=LT, metadata={op_name="compare.35"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.36 = s32[]{:T(128)S(6)} xor(%constant.1386, %bitcast-convert.35), metadata={op_name="xor.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.118 = s32[]{:T(128)S(6)} select(%compare.128, %xor.36, %bitcast-convert.35), metadata={op_name="select.14"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %convert.386 = f32[]{:T(128)S(6)} convert(%top_k.6), metadata={op_name="convert.17"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %bitcast-convert.36 = s32[]{:T(128)S(6)} bitcast-convert(%convert.386), metadata={op_name="bitcast-convert.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.129 = pred[]{:T(512)S(6)} compare(%bitcast-convert.36, %constant.1369), direction=LT, metadata={op_name="compare.36"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.37 = s32[]{:T(128)S(6)} xor(%constant.1370, %bitcast-convert.36), metadata={op_name="xor.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.129 = pred[]{:T(512)S(6)} compare(%bitcast-convert.36, %constant.1385), direction=LT, metadata={op_name="compare.36"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.37 = s32[]{:T(128)S(6)} xor(%constant.1386, %bitcast-convert.36), metadata={op_name="xor.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.119 = s32[]{:T(128)S(6)} select(%compare.129, %xor.37, %bitcast-convert.36), metadata={op_name="select.15"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %compare.130 = pred[]{:T(512)S(6)} compare(%select.118, %select.119), direction=GT, metadata={op_name="compare.1"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %compare.131 = pred[]{:T(512)S(6)} compare(%select.119, %select.118), direction=GT, metadata={op_name="compare.108"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} @@ -420,12 +420,12 @@ StackFrames %param_0.68 = s32[256]{0:T(256)S(1)} parameter(0) %param_1.114 = s32[1024]{0:T(1024)S(1)} parameter(1) %custom-call.19 = s32[1024]{0:T(1024)} custom-call(%param_1.114), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %slice.926 = s32[263]{0:T(512)} slice(%custom-call.19), slice={[0:263]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3634 = s32[263]{0:T(512)} reshape(%slice.926), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.1037 = s32[263]{0:T(512)} transpose(%reshape.3634), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.213 = s32[263]{0:T(512)} gather(%param_0.68, %transpose.1037), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.1036 = s32[263]{0:T(512)} transpose(%gather.213), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3633 = s32[263]{0:T(512)S(1)} reshape(%transpose.1036), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %slice.898 = s32[263]{0:T(512)} slice(%custom-call.19), slice={[0:263]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.3614 = s32[263]{0:T(512)} reshape(%slice.898), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.1037 = s32[263]{0:T(512)} transpose(%reshape.3614), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.209 = s32[263]{0:T(512)} gather(%param_0.68, %transpose.1037), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.1036 = s32[263]{0:T(512)} transpose(%gather.209), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.3613 = s32[263]{0:T(512)S(1)} reshape(%transpose.1036), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} } %region_27.34.clone.1 (reduce-window.350: s32[], reduce-window.351: s32[]) -> s32[] { @@ -464,12 +464,12 @@ StackFrames %param_0.71 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.116 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.21 = s32[4096]{0:T(1024)} custom-call(%param_1.116), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.928 = s32[4096]{0:T(1024)} slice(%custom-call.21), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3657 = s32[4096]{0:T(1024)} reshape(%slice.928), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.1043 = s32[4096]{0:T(1024)} transpose(%reshape.3657), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.214 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.71, %transpose.1043), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.1042 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.214), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3656 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1042), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.900 = s32[4096]{0:T(1024)} slice(%custom-call.21), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3637 = s32[4096]{0:T(1024)} reshape(%slice.900), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.1043 = s32[4096]{0:T(1024)} transpose(%reshape.3637), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.210 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.71, %transpose.1043), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.1042 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.210), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3636 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1042), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %region_31.39 (sort.50: s32[], sort.51: s32[], sort.52: s32[], sort.53: s32[], sort.54: s32[], sort.55: s32[]) -> pred[] { @@ -490,12 +490,12 @@ StackFrames %param_0.72 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) %param_1.118 = s32[4096]{0:T(1024)S(1)} parameter(1) %custom-call.23 = s32[4096]{0:T(1024)} custom-call(%param_1.118), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.930 = s32[4096]{0:T(1024)} slice(%custom-call.23), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3659 = s32[4096]{0:T(1024)} reshape(%slice.930), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.1045 = s32[4096]{0:T(1024)} transpose(%reshape.3659), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.215 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.72, %transpose.1045), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.1044 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.215), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3658 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1044), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.902 = s32[4096]{0:T(1024)} slice(%custom-call.23), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.3639 = s32[4096]{0:T(1024)} reshape(%slice.902), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.1045 = s32[4096]{0:T(1024)} transpose(%reshape.3639), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.211 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.72, %transpose.1045), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.1044 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.211), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.3638 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1044), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} } %compare (name: s32[], name.1: s32[], name.2: bf16[], name.3: bf16[]) -> pred[] { @@ -503,7 +503,7 @@ StackFrames %name.3 = bf16[] parameter(3) %name = s32[] parameter(0) %name.1 = s32[] parameter(1) - ROOT %compare.385 = pred[] compare(%name, %name.1), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.377 = pred[] compare(%name, %name.1), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %compare.1 (name.4: s32[], name.5: s32[], name.6: f32[], name.7: f32[]) -> pred[] { @@ -511,7 +511,7 @@ StackFrames %name.7 = f32[] parameter(3) %name.4 = s32[] parameter(0) %name.5 = s32[] parameter(1) - ROOT %compare.386 = pred[] compare(%name.4, %name.5), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.378 = pred[] compare(%name.4, %name.5), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %compare.2 (name.8: s32[], name.9: s32[], name.10: f32[], name.11: f32[]) -> pred[] { @@ -519,7 +519,7 @@ StackFrames %name.11 = f32[] parameter(3) %name.8 = s32[] parameter(0) %name.9 = s32[] parameter(1) - ROOT %compare.387 = pred[] compare(%name.8, %name.9), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.379 = pred[] compare(%name.8, %name.9), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %compare.3 (name.12: s32[], name.13: s32[], name.14: f32[], name.15: f32[]) -> pred[] { @@ -527,7 +527,7 @@ StackFrames %name.15 = f32[] parameter(3) %name.12 = s32[] parameter(0) %name.13 = s32[] parameter(1) - ROOT %compare.388 = pred[] compare(%name.12, %name.13), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.380 = pred[] compare(%name.12, %name.13), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %compare.4 (name.16: s32[], name.17: s32[], name.18: f32[], name.19: f32[]) -> pred[] { @@ -535,52 +535,52 @@ StackFrames %name.19 = f32[] parameter(3) %name.16 = s32[] parameter(0) %name.17 = s32[] parameter(1) - ROOT %compare.389 = pred[] compare(%name.16, %name.17), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %compare.381 = pred[] compare(%name.16, %name.17), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%called_computation.13 (param_0.4523: s32[256]) -> s32[256] { - %param_0.4523 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2073 = s32[256]{0:T(256)} copy(%param_0.4523), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"1134","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.13 (param_0.4524: s32[256]) -> s32[256] { + %param_0.4524 = s32[256]{0:T(256)} parameter(0) + ROOT %copy.2073 = s32[256]{0:T(256)} copy(%param_0.4524), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"1134","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.13 (param_0.4524: s32[256]) -> s32[256] { - %param_0.4524 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2074.cloned.1 = s32[256]{0:T(256)} call(%param_0.4524), to_apply=%called_computation.13 +%async_computation.13 (param_0.4525: s32[256]) -> s32[256] { + %param_0.4525 = s32[256]{0:T(256)} parameter(0) + ROOT %copy.2074.cloned.1 = s32[256]{0:T(256)} call(%param_0.4525), to_apply=%called_computation.13 }, execution_thread="sparsecore" %region_49.59 (scatter-add.14: s32[], scatter-add.15: s32[]) -> s32[] { %scatter-add.14 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.15 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1352 = s32[]{:T(128)S(7)} add(%scatter-add.14, %scatter-add.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.22.clone.clone (param_0.4525: s32[256], param_1.5325: s32[4096], param_2.4494: s32[4096]) -> s32[256] { - %param_0.4525 = s32[256]{0:T(256)} parameter(0) - %param_1.5325 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3923 = s32[4096]{0:T(1024)} reshape(%param_1.5325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} - %transpose.1100 = s32[4096]{0:T(1024)} transpose(%reshape.3923), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} - %param_2.4494 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3924 = s32[4096]{0:T(1024)} reshape(%param_2.4494), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - %transpose.1101 = s32[4096]{0:T(1024)} transpose(%reshape.3924), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.237 = s32[256]{0:T(256)} scatter(%param_0.4525, %transpose.1100, %transpose.1101), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_49.59, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %add.1362 = s32[]{:T(128)S(7)} add(%scatter-add.14, %scatter-add.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.14 (param_0.4526: s32[256], param_1.5326: s32[4096], param_2.4495: s32[4096]) -> s32[256] { +%fused_computation.22.clone.clone (param_0.4526: s32[256], param_1.5321: s32[4096], param_2.4492: s32[4096]) -> s32[256] { %param_0.4526 = s32[256]{0:T(256)} parameter(0) - %param_1.5326 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4495 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4526, %param_1.5326, %param_2.4495), kind=kCustom, calls=%fused_computation.22.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5321 = s32[4096]{0:T(1024)} parameter(1) + %reshape.3903 = s32[4096]{0:T(1024)} reshape(%param_1.5321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} + %transpose.1100 = s32[4096]{0:T(1024)} transpose(%reshape.3903), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} + %param_2.4492 = s32[4096]{0:T(1024)} parameter(2) + %reshape.3904 = s32[4096]{0:T(1024)} reshape(%param_2.4492), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + %transpose.1101 = s32[4096]{0:T(1024)} transpose(%reshape.3904), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.231 = s32[256]{0:T(256)} scatter(%param_0.4526, %transpose.1100, %transpose.1101), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_49.59, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.14 (param_0.4527: s32[256], param_1.5327: s32[4096], param_2.4496: s32[4096]) -> s32[256] { +%called_computation.14 (param_0.4527: s32[256], param_1.5322: s32[4096], param_2.4493: s32[4096]) -> s32[256] { %param_0.4527 = s32[256]{0:T(256)} parameter(0) - %param_1.5327 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4496 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4527, %param_1.5327, %param_2.4496), to_apply=%called_computation.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} + %param_1.5322 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4493 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4527, %param_1.5322, %param_2.4493), kind=kCustom, calls=%fused_computation.22.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation (param_0.84: s32[256], param_1.136: s32[4096], param_2.80: s32[4096], param_3.3085: token[]) -> s32[256] { - %param_3.3085 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.14 (param_0.4528: s32[256], param_1.5323: s32[4096], param_2.4494: s32[4096]) -> s32[256] { + %param_0.4528 = s32[256]{0:T(256)} parameter(0) + %param_1.5323 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4494 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4528, %param_1.5323, %param_2.4494), to_apply=%called_computation.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation (param_0.84: s32[256], param_1.136: s32[4096], param_2.80: s32[4096], param_3.3083: token[]) -> s32[256] { + %param_3.3083 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.84 = s32[256]{0:T(256)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.136 = s32[4096]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.80 = s32[4096]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -590,57 +590,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.40.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.40.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation (param_0.85: s32[256], param_1.137: s32[4096], param_2.81: s32[4096], param_3.3084: token[]) -> s32[256] { - %param_3.3084 = token[] parameter(3) +%async_computation (param_0.85: s32[256], param_1.137: s32[4096], param_2.81: s32[4096], param_3.3082: token[]) -> s32[256] { + %param_3.3082 = token[] parameter(3) %param_0.85 = s32[256]{0:T(256)} parameter(0) %param_1.137 = s32[4096]{0:T(1024)} parameter(1) %param_2.81 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.2.cloned.1 = s32[256]{0:T(256)} call(%param_0.85, %param_1.137, %param_2.81, %param_3.3084), to_apply=%called_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.2.cloned.1 = s32[256]{0:T(256)} call(%param_0.85, %param_1.137, %param_2.81, %param_3.3082), to_apply=%called_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.15 (param_0.4528: f32[9]) -> f32[9] { - %param_0.4528 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2075 = f32[9]{0:T(128)} copy(%param_0.4528), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.15 (param_0.4529: f32[9]) -> f32[9] { + %param_0.4529 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2075 = f32[9]{0:T(128)} copy(%param_0.4529), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.15 (param_0.4529: f32[9]) -> f32[9] { - %param_0.4529 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2076.cloned.1 = f32[9]{0:T(128)} call(%param_0.4529), to_apply=%called_computation.15 +%async_computation.15 (param_0.4530: f32[9]) -> f32[9] { + %param_0.4530 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2076.cloned.1 = f32[9]{0:T(128)} call(%param_0.4530), to_apply=%called_computation.15 }, execution_thread="sparsecore" %region_61.72 (scatter-add.24: f32[], scatter-add.25: f32[]) -> f32[] { %scatter-add.24 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.25 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1358 = f32[]{:T(128)S(7)} add(%scatter-add.24, %scatter-add.25), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1368 = f32[]{:T(128)S(7)} add(%scatter-add.24, %scatter-add.25), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.24.clone.clone (param_0.4530: f32[9], param_1.5328: s32[256], param_2.4497: f32[256]) -> f32[9] { - %param_0.4530 = f32[9]{0:T(128)} parameter(0) - %param_1.5328 = s32[256]{0:T(256)} parameter(1) - %reshape.3925 = s32[256]{0:T(256)} reshape(%param_1.5328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1102 = s32[256]{0:T(256)} transpose(%reshape.3925), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4497 = f32[256]{0:T(256)} parameter(2) - %reshape.3926 = f32[256]{0:T(256)} reshape(%param_2.4497), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1103 = f32[256]{0:T(256)} transpose(%reshape.3926), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.238 = f32[9]{0:T(128)} scatter(%param_0.4530, %transpose.1102, %transpose.1103), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_61.72, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.16 (param_0.4531: f32[9], param_1.5329: s32[256], param_2.4498: f32[256]) -> f32[9] { +%fused_computation.24.clone.clone (param_0.4531: f32[9], param_1.5324: s32[256], param_2.4495: f32[256]) -> f32[9] { %param_0.4531 = f32[9]{0:T(128)} parameter(0) - %param_1.5329 = s32[256]{0:T(256)} parameter(1) - %param_2.4498 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.41 = f32[9]{0:T(128)} fusion(%param_0.4531, %param_1.5329, %param_2.4498), kind=kCustom, calls=%fused_computation.24.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5324 = s32[256]{0:T(256)} parameter(1) + %reshape.3905 = s32[256]{0:T(256)} reshape(%param_1.5324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1102 = s32[256]{0:T(256)} transpose(%reshape.3905), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4495 = f32[256]{0:T(256)} parameter(2) + %reshape.3906 = f32[256]{0:T(256)} reshape(%param_2.4495), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1103 = f32[256]{0:T(256)} transpose(%reshape.3906), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.232 = f32[9]{0:T(128)} scatter(%param_0.4531, %transpose.1102, %transpose.1103), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_61.72, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.16 (param_0.4532: f32[9], param_1.5330: s32[256], param_2.4499: f32[256]) -> f32[9] { +%called_computation.16 (param_0.4532: f32[9], param_1.5325: s32[256], param_2.4496: f32[256]) -> f32[9] { %param_0.4532 = f32[9]{0:T(128)} parameter(0) - %param_1.5330 = s32[256]{0:T(256)} parameter(1) - %param_2.4499 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.42.cloned.1 = f32[9]{0:T(128)} call(%param_0.4532, %param_1.5330, %param_2.4499), to_apply=%called_computation.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5325 = s32[256]{0:T(256)} parameter(1) + %param_2.4496 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.41 = f32[9]{0:T(128)} fusion(%param_0.4532, %param_1.5325, %param_2.4496), kind=kCustom, calls=%fused_computation.24.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.1 (param_0.87: f32[9], param_1.139: s32[256], param_2.83: f32[256], param_3.3099: token[]) -> f32[9] { - %param_3.3099 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.16 (param_0.4533: f32[9], param_1.5326: s32[256], param_2.4497: f32[256]) -> f32[9] { + %param_0.4533 = f32[9]{0:T(128)} parameter(0) + %param_1.5326 = s32[256]{0:T(256)} parameter(1) + %param_2.4497 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.42.cloned.1 = f32[9]{0:T(128)} call(%param_0.4533, %param_1.5326, %param_2.4497), to_apply=%called_computation.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.1 (param_0.87: f32[9], param_1.139: s32[256], param_2.83: f32[256], param_3.3097: token[]) -> f32[9] { + %param_3.3097 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.87 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.139 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.83 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -650,57 +650,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.42.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.42.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.1 (param_0.88: f32[9], param_1.140: s32[256], param_2.84: f32[256], param_3.3098: token[]) -> f32[9] { - %param_3.3098 = token[] parameter(3) +%async_computation.1 (param_0.88: f32[9], param_1.140: s32[256], param_2.84: f32[256], param_3.3096: token[]) -> f32[9] { + %param_3.3096 = token[] parameter(3) %param_0.88 = f32[9]{0:T(128)} parameter(0) %param_1.140 = s32[256]{0:T(256)} parameter(1) %param_2.84 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.5.cloned.1 = f32[9]{0:T(128)} call(%param_0.88, %param_1.140, %param_2.84, %param_3.3098), to_apply=%called_computation.1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.5.cloned.1 = f32[9]{0:T(128)} call(%param_0.88, %param_1.140, %param_2.84, %param_3.3096), to_apply=%called_computation.1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.17 (param_0.4533: s32[263]) -> s32[263] { - %param_0.4533 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2077 = s32[263]{0:T(512)} copy(%param_0.4533), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.17 (param_0.4534: s32[263]) -> s32[263] { + %param_0.4534 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2077 = s32[263]{0:T(512)} copy(%param_0.4534), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.17 (param_0.4534: s32[263]) -> s32[263] { - %param_0.4534 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2078.cloned.1 = s32[263]{0:T(512)} call(%param_0.4534), to_apply=%called_computation.17 +%async_computation.17 (param_0.4535: s32[263]) -> s32[263] { + %param_0.4535 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2078.cloned.1 = s32[263]{0:T(512)} call(%param_0.4535), to_apply=%called_computation.17 }, execution_thread="sparsecore" %region_63.74 (scatter-add.28: s32[], scatter-add.29: s32[]) -> s32[] { %scatter-add.28 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.29 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1359 = s32[]{:T(128)S(7)} add(%scatter-add.28, %scatter-add.29), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1369 = s32[]{:T(128)S(7)} add(%scatter-add.28, %scatter-add.29), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.25.clone.clone (param_0.4535: s32[263], param_1.5331: s32[8], param_2.4500: s32[8]) -> s32[263] { - %param_0.4535 = s32[263]{0:T(512)} parameter(0) - %param_1.5331 = s32[8]{0:T(128)} parameter(1) - %reshape.3927 = s32[8]{0:T(128)} reshape(%param_1.5331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1104 = s32[8]{0:T(128)} transpose(%reshape.3927), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4500 = s32[8]{0:T(128)} parameter(2) - %reshape.3928 = s32[8]{0:T(128)} reshape(%param_2.4500), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1105 = s32[8]{0:T(128)} transpose(%reshape.3928), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.239 = s32[263]{0:T(512)} scatter(%param_0.4535, %transpose.1104, %transpose.1105), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_63.74, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.18 (param_0.4536: s32[263], param_1.5332: s32[8], param_2.4501: s32[8]) -> s32[263] { +%fused_computation.25.clone.clone (param_0.4536: s32[263], param_1.5327: s32[8], param_2.4498: s32[8]) -> s32[263] { %param_0.4536 = s32[263]{0:T(512)} parameter(0) - %param_1.5332 = s32[8]{0:T(128)} parameter(1) - %param_2.4501 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.43 = s32[263]{0:T(512)} fusion(%param_0.4536, %param_1.5332, %param_2.4501), kind=kCustom, calls=%fused_computation.25.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5327 = s32[8]{0:T(128)} parameter(1) + %reshape.3907 = s32[8]{0:T(128)} reshape(%param_1.5327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1104 = s32[8]{0:T(128)} transpose(%reshape.3907), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4498 = s32[8]{0:T(128)} parameter(2) + %reshape.3908 = s32[8]{0:T(128)} reshape(%param_2.4498), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %transpose.1105 = s32[8]{0:T(128)} transpose(%reshape.3908), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + ROOT %scatter-add.233 = s32[263]{0:T(512)} scatter(%param_0.4536, %transpose.1104, %transpose.1105), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_63.74, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.18 (param_0.4537: s32[263], param_1.5333: s32[8], param_2.4502: s32[8]) -> s32[263] { +%called_computation.18 (param_0.4537: s32[263], param_1.5328: s32[8], param_2.4499: s32[8]) -> s32[263] { %param_0.4537 = s32[263]{0:T(512)} parameter(0) - %param_1.5333 = s32[8]{0:T(128)} parameter(1) - %param_2.4502 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.44.cloned.1 = s32[263]{0:T(512)} call(%param_0.4537, %param_1.5333, %param_2.4502), to_apply=%called_computation.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5328 = s32[8]{0:T(128)} parameter(1) + %param_2.4499 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.43 = s32[263]{0:T(512)} fusion(%param_0.4537, %param_1.5328, %param_2.4499), kind=kCustom, calls=%fused_computation.25.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.2 (param_0.90: s32[263], param_1.142: s32[8], param_2.86: s32[8], param_3.3105: token[]) -> s32[263] { - %param_3.3105 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.18 (param_0.4538: s32[263], param_1.5329: s32[8], param_2.4500: s32[8]) -> s32[263] { + %param_0.4538 = s32[263]{0:T(512)} parameter(0) + %param_1.5329 = s32[8]{0:T(128)} parameter(1) + %param_2.4500 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.44.cloned.1 = s32[263]{0:T(512)} call(%param_0.4538, %param_1.5329, %param_2.4500), to_apply=%called_computation.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.2 (param_0.90: s32[263], param_1.142: s32[8], param_2.86: s32[8], param_3.3103: token[]) -> s32[263] { + %param_3.3103 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.90 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.142 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.86 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -710,57 +710,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.44.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.44.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.2 (param_0.91: s32[263], param_1.143: s32[8], param_2.87: s32[8], param_3.3104: token[]) -> s32[263] { - %param_3.3104 = token[] parameter(3) +%async_computation.2 (param_0.91: s32[263], param_1.143: s32[8], param_2.87: s32[8], param_3.3102: token[]) -> s32[263] { + %param_3.3102 = token[] parameter(3) %param_0.91 = s32[263]{0:T(512)} parameter(0) %param_1.143 = s32[8]{0:T(128)} parameter(1) %param_2.87 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.8.cloned.1 = s32[263]{0:T(512)} call(%param_0.91, %param_1.143, %param_2.87, %param_3.3104), to_apply=%called_computation.2, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.8.cloned.1 = s32[263]{0:T(512)} call(%param_0.91, %param_1.143, %param_2.87, %param_3.3102), to_apply=%called_computation.2, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.19 (param_0.4538: s32[263]) -> s32[263] { - %param_0.4538 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2079 = s32[263]{0:T(512)} copy(%param_0.4538), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.19 (param_0.4539: s32[263]) -> s32[263] { + %param_0.4539 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2079 = s32[263]{0:T(512)} copy(%param_0.4539), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.19 (param_0.4539: s32[263]) -> s32[263] { - %param_0.4539 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2080.cloned.1 = s32[263]{0:T(512)} call(%param_0.4539), to_apply=%called_computation.19 +%async_computation.19 (param_0.4540: s32[263]) -> s32[263] { + %param_0.4540 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2080.cloned.1 = s32[263]{0:T(512)} call(%param_0.4540), to_apply=%called_computation.19 }, execution_thread="sparsecore" %region_73.86.clone (scatter-add.163: s32[], scatter-add.164: s32[]) -> s32[] { %scatter-add.163 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.164 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2474 = s32[]{:T(128)S(7)} add(%scatter-add.163, %scatter-add.164), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.2482 = s32[]{:T(128)S(7)} add(%scatter-add.163, %scatter-add.164), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.26.clone.clone (param_0.4540: s32[263], param_1.5334: s32[256], param_2.4503: s32[256]) -> s32[263] { - %param_0.4540 = s32[263]{0:T(512)} parameter(0) - %param_1.5334 = s32[256]{0:T(256)} parameter(1) - %reshape.3929 = s32[256]{0:T(256)} reshape(%param_1.5334), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1106 = s32[256]{0:T(256)} transpose(%reshape.3929), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4503 = s32[256]{0:T(256)} parameter(2) - %reshape.3930 = s32[256]{0:T(256)} reshape(%param_2.4503), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1107 = s32[256]{0:T(256)} transpose(%reshape.3930), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.240 = s32[263]{0:T(512)} scatter(%param_0.4540, %transpose.1106, %transpose.1107), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_73.86.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.20 (param_0.4541: s32[263], param_1.5335: s32[256], param_2.4504: s32[256]) -> s32[263] { +%fused_computation.26.clone.clone (param_0.4541: s32[263], param_1.5330: s32[256], param_2.4501: s32[256]) -> s32[263] { %param_0.4541 = s32[263]{0:T(512)} parameter(0) - %param_1.5335 = s32[256]{0:T(256)} parameter(1) - %param_2.4504 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.45 = s32[263]{0:T(512)} fusion(%param_0.4541, %param_1.5335, %param_2.4504), kind=kCustom, calls=%fused_computation.26.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5330 = s32[256]{0:T(256)} parameter(1) + %reshape.3909 = s32[256]{0:T(256)} reshape(%param_1.5330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1106 = s32[256]{0:T(256)} transpose(%reshape.3909), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4501 = s32[256]{0:T(256)} parameter(2) + %reshape.3910 = s32[256]{0:T(256)} reshape(%param_2.4501), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1107 = s32[256]{0:T(256)} transpose(%reshape.3910), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.234 = s32[263]{0:T(512)} scatter(%param_0.4541, %transpose.1106, %transpose.1107), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_73.86.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.20 (param_0.4542: s32[263], param_1.5336: s32[256], param_2.4505: s32[256]) -> s32[263] { +%called_computation.20 (param_0.4542: s32[263], param_1.5331: s32[256], param_2.4502: s32[256]) -> s32[263] { %param_0.4542 = s32[263]{0:T(512)} parameter(0) - %param_1.5336 = s32[256]{0:T(256)} parameter(1) - %param_2.4505 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.46.cloned.1 = s32[263]{0:T(512)} call(%param_0.4542, %param_1.5336, %param_2.4505), to_apply=%called_computation.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5331 = s32[256]{0:T(256)} parameter(1) + %param_2.4502 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.45 = s32[263]{0:T(512)} fusion(%param_0.4542, %param_1.5331, %param_2.4502), kind=kCustom, calls=%fused_computation.26.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.3 (param_0.93: s32[263], param_1.145: s32[256], param_2.89: s32[256], param_3.3091: token[]) -> s32[263] { - %param_3.3091 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.20 (param_0.4543: s32[263], param_1.5332: s32[256], param_2.4503: s32[256]) -> s32[263] { + %param_0.4543 = s32[263]{0:T(512)} parameter(0) + %param_1.5332 = s32[256]{0:T(256)} parameter(1) + %param_2.4503 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.46.cloned.1 = s32[263]{0:T(512)} call(%param_0.4543, %param_1.5332, %param_2.4503), to_apply=%called_computation.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.3 (param_0.93: s32[263], param_1.145: s32[256], param_2.89: s32[256], param_3.3089: token[]) -> s32[263] { + %param_3.3089 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.93 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.145 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.89 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -770,57 +770,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.46.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.46.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.3 (param_0.94: s32[263], param_1.146: s32[256], param_2.90: s32[256], param_3.3090: token[]) -> s32[263] { - %param_3.3090 = token[] parameter(3) +%async_computation.3 (param_0.94: s32[263], param_1.146: s32[256], param_2.90: s32[256], param_3.3088: token[]) -> s32[263] { + %param_3.3088 = token[] parameter(3) %param_0.94 = s32[263]{0:T(512)} parameter(0) %param_1.146 = s32[256]{0:T(256)} parameter(1) %param_2.90 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.11.cloned.1 = s32[263]{0:T(512)} call(%param_0.94, %param_1.146, %param_2.90, %param_3.3090), to_apply=%called_computation.3, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.11.cloned.1 = s32[263]{0:T(512)} call(%param_0.94, %param_1.146, %param_2.90, %param_3.3088), to_apply=%called_computation.3, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.21 (param_0.4543: f32[9]) -> f32[9] { - %param_0.4543 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2081 = f32[9]{0:T(128)} copy(%param_0.4543), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.21 (param_0.4544: f32[9]) -> f32[9] { + %param_0.4544 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2081 = f32[9]{0:T(128)} copy(%param_0.4544), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.21 (param_0.4544: f32[9]) -> f32[9] { - %param_0.4544 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2082.cloned.1 = f32[9]{0:T(128)} call(%param_0.4544), to_apply=%called_computation.21 +%async_computation.21 (param_0.4545: f32[9]) -> f32[9] { + %param_0.4545 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2082.cloned.1 = f32[9]{0:T(128)} call(%param_0.4545), to_apply=%called_computation.21 }, execution_thread="sparsecore" %region_79.95.clone (scatter-add.167: f32[], scatter-add.168: f32[]) -> f32[] { %scatter-add.167 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.168 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2476 = f32[]{:T(128)S(7)} add(%scatter-add.167, %scatter-add.168), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.27.clone.clone (param_0.4545: f32[9], param_1.5337: s32[256], param_2.4506: f32[256]) -> f32[9] { - %param_0.4545 = f32[9]{0:T(128)} parameter(0) - %param_1.5337 = s32[256]{0:T(256)} parameter(1) - %reshape.3931 = s32[256]{0:T(256)} reshape(%param_1.5337), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1108 = s32[256]{0:T(256)} transpose(%reshape.3931), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4506 = f32[256]{0:T(256)} parameter(2) - %reshape.3932 = f32[256]{0:T(256)} reshape(%param_2.4506), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1109 = f32[256]{0:T(256)} transpose(%reshape.3932), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.241 = f32[9]{0:T(128)} scatter(%param_0.4545, %transpose.1108, %transpose.1109), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_79.95.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %add.2484 = f32[]{:T(128)S(7)} add(%scatter-add.167, %scatter-add.168), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.22 (param_0.4546: f32[9], param_1.5338: s32[256], param_2.4507: f32[256]) -> f32[9] { +%fused_computation.27.clone.clone (param_0.4546: f32[9], param_1.5333: s32[256], param_2.4504: f32[256]) -> f32[9] { %param_0.4546 = f32[9]{0:T(128)} parameter(0) - %param_1.5338 = s32[256]{0:T(256)} parameter(1) - %param_2.4507 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.47 = f32[9]{0:T(128)} fusion(%param_0.4546, %param_1.5338, %param_2.4507), kind=kCustom, calls=%fused_computation.27.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5333 = s32[256]{0:T(256)} parameter(1) + %reshape.3911 = s32[256]{0:T(256)} reshape(%param_1.5333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1108 = s32[256]{0:T(256)} transpose(%reshape.3911), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4504 = f32[256]{0:T(256)} parameter(2) + %reshape.3912 = f32[256]{0:T(256)} reshape(%param_2.4504), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1109 = f32[256]{0:T(256)} transpose(%reshape.3912), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.235 = f32[9]{0:T(128)} scatter(%param_0.4546, %transpose.1108, %transpose.1109), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_79.95.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.22 (param_0.4547: f32[9], param_1.5339: s32[256], param_2.4508: f32[256]) -> f32[9] { +%called_computation.22 (param_0.4547: f32[9], param_1.5334: s32[256], param_2.4505: f32[256]) -> f32[9] { %param_0.4547 = f32[9]{0:T(128)} parameter(0) - %param_1.5339 = s32[256]{0:T(256)} parameter(1) - %param_2.4508 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.48.cloned.1 = f32[9]{0:T(128)} call(%param_0.4547, %param_1.5339, %param_2.4508), to_apply=%called_computation.22, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5334 = s32[256]{0:T(256)} parameter(1) + %param_2.4505 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.47 = f32[9]{0:T(128)} fusion(%param_0.4547, %param_1.5334, %param_2.4505), kind=kCustom, calls=%fused_computation.27.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.4 (param_0.96: f32[9], param_1.148: s32[256], param_2.92: f32[256], param_3.3097: token[]) -> f32[9] { - %param_3.3097 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.22 (param_0.4548: f32[9], param_1.5335: s32[256], param_2.4506: f32[256]) -> f32[9] { + %param_0.4548 = f32[9]{0:T(128)} parameter(0) + %param_1.5335 = s32[256]{0:T(256)} parameter(1) + %param_2.4506 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.48.cloned.1 = f32[9]{0:T(128)} call(%param_0.4548, %param_1.5335, %param_2.4506), to_apply=%called_computation.22, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.4 (param_0.96: f32[9], param_1.148: s32[256], param_2.92: f32[256], param_3.3095: token[]) -> f32[9] { + %param_3.3095 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.96 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.148 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.92 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -830,57 +830,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.48.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.48.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.4 (param_0.97: f32[9], param_1.149: s32[256], param_2.93: f32[256], param_3.3096: token[]) -> f32[9] { - %param_3.3096 = token[] parameter(3) +%async_computation.4 (param_0.97: f32[9], param_1.149: s32[256], param_2.93: f32[256], param_3.3094: token[]) -> f32[9] { + %param_3.3094 = token[] parameter(3) %param_0.97 = f32[9]{0:T(128)} parameter(0) %param_1.149 = s32[256]{0:T(256)} parameter(1) %param_2.93 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.14.cloned.1 = f32[9]{0:T(128)} call(%param_0.97, %param_1.149, %param_2.93, %param_3.3096), to_apply=%called_computation.4, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.14.cloned.1 = f32[9]{0:T(128)} call(%param_0.97, %param_1.149, %param_2.93, %param_3.3094), to_apply=%called_computation.4, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.23 (param_0.4548: s32[263]) -> s32[263] { - %param_0.4548 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2083 = s32[263]{0:T(512)} copy(%param_0.4548), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.23 (param_0.4549: s32[263]) -> s32[263] { + %param_0.4549 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2083 = s32[263]{0:T(512)} copy(%param_0.4549), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.23 (param_0.4549: s32[263]) -> s32[263] { - %param_0.4549 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2084.cloned.1 = s32[263]{0:T(512)} call(%param_0.4549), to_apply=%called_computation.23 +%async_computation.23 (param_0.4550: s32[263]) -> s32[263] { + %param_0.4550 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2084.cloned.1 = s32[263]{0:T(512)} call(%param_0.4550), to_apply=%called_computation.23 }, execution_thread="sparsecore" %region_81.97.clone (scatter-add.171: s32[], scatter-add.172: s32[]) -> s32[] { %scatter-add.171 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.172 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2478 = s32[]{:T(128)S(7)} add(%scatter-add.171, %scatter-add.172), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.2486 = s32[]{:T(128)S(7)} add(%scatter-add.171, %scatter-add.172), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.28.clone.clone (param_0.4550: s32[263], param_1.5340: s32[8], param_2.4509: s32[8]) -> s32[263] { - %param_0.4550 = s32[263]{0:T(512)} parameter(0) - %param_1.5340 = s32[8]{0:T(128)} parameter(1) - %reshape.3933 = s32[8]{0:T(128)} reshape(%param_1.5340), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1110 = s32[8]{0:T(128)} transpose(%reshape.3933), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4509 = s32[8]{0:T(128)} parameter(2) - %reshape.3934 = s32[8]{0:T(128)} reshape(%param_2.4509), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1111 = s32[8]{0:T(128)} transpose(%reshape.3934), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.242 = s32[263]{0:T(512)} scatter(%param_0.4550, %transpose.1110, %transpose.1111), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_81.97.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.24 (param_0.4551: s32[263], param_1.5341: s32[8], param_2.4510: s32[8]) -> s32[263] { +%fused_computation.28.clone.clone (param_0.4551: s32[263], param_1.5336: s32[8], param_2.4507: s32[8]) -> s32[263] { %param_0.4551 = s32[263]{0:T(512)} parameter(0) - %param_1.5341 = s32[8]{0:T(128)} parameter(1) - %param_2.4510 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.49 = s32[263]{0:T(512)} fusion(%param_0.4551, %param_1.5341, %param_2.4510), kind=kCustom, calls=%fused_computation.28.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5336 = s32[8]{0:T(128)} parameter(1) + %reshape.3913 = s32[8]{0:T(128)} reshape(%param_1.5336), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1110 = s32[8]{0:T(128)} transpose(%reshape.3913), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4507 = s32[8]{0:T(128)} parameter(2) + %reshape.3914 = s32[8]{0:T(128)} reshape(%param_2.4507), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %transpose.1111 = s32[8]{0:T(128)} transpose(%reshape.3914), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + ROOT %scatter-add.236 = s32[263]{0:T(512)} scatter(%param_0.4551, %transpose.1110, %transpose.1111), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_81.97.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.24 (param_0.4552: s32[263], param_1.5342: s32[8], param_2.4511: s32[8]) -> s32[263] { +%called_computation.24 (param_0.4552: s32[263], param_1.5337: s32[8], param_2.4508: s32[8]) -> s32[263] { %param_0.4552 = s32[263]{0:T(512)} parameter(0) - %param_1.5342 = s32[8]{0:T(128)} parameter(1) - %param_2.4511 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.50.cloned.1 = s32[263]{0:T(512)} call(%param_0.4552, %param_1.5342, %param_2.4511), to_apply=%called_computation.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5337 = s32[8]{0:T(128)} parameter(1) + %param_2.4508 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.49 = s32[263]{0:T(512)} fusion(%param_0.4552, %param_1.5337, %param_2.4508), kind=kCustom, calls=%fused_computation.28.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.5 (param_0.99: s32[263], param_1.151: s32[8], param_2.95: s32[8], param_3.3107: token[]) -> s32[263] { - %param_3.3107 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.24 (param_0.4553: s32[263], param_1.5338: s32[8], param_2.4509: s32[8]) -> s32[263] { + %param_0.4553 = s32[263]{0:T(512)} parameter(0) + %param_1.5338 = s32[8]{0:T(128)} parameter(1) + %param_2.4509 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.50.cloned.1 = s32[263]{0:T(512)} call(%param_0.4553, %param_1.5338, %param_2.4509), to_apply=%called_computation.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.5 (param_0.99: s32[263], param_1.151: s32[8], param_2.95: s32[8], param_3.3105: token[]) -> s32[263] { + %param_3.3105 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.99 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.151 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.95 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -890,57 +890,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.50.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.50.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.5 (param_0.100: s32[263], param_1.152: s32[8], param_2.96: s32[8], param_3.3106: token[]) -> s32[263] { - %param_3.3106 = token[] parameter(3) +%async_computation.5 (param_0.100: s32[263], param_1.152: s32[8], param_2.96: s32[8], param_3.3104: token[]) -> s32[263] { + %param_3.3104 = token[] parameter(3) %param_0.100 = s32[263]{0:T(512)} parameter(0) %param_1.152 = s32[8]{0:T(128)} parameter(1) %param_2.96 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.17.cloned.1 = s32[263]{0:T(512)} call(%param_0.100, %param_1.152, %param_2.96, %param_3.3106), to_apply=%called_computation.5, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.17.cloned.1 = s32[263]{0:T(512)} call(%param_0.100, %param_1.152, %param_2.96, %param_3.3104), to_apply=%called_computation.5, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.25 (param_0.4553: s32[263]) -> s32[263] { - %param_0.4553 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2085 = s32[263]{0:T(512)} copy(%param_0.4553), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.25 (param_0.4554: s32[263]) -> s32[263] { + %param_0.4554 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2085 = s32[263]{0:T(512)} copy(%param_0.4554), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.25 (param_0.4554: s32[263]) -> s32[263] { - %param_0.4554 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2086.cloned.1 = s32[263]{0:T(512)} call(%param_0.4554), to_apply=%called_computation.25 +%async_computation.25 (param_0.4555: s32[263]) -> s32[263] { + %param_0.4555 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2086.cloned.1 = s32[263]{0:T(512)} call(%param_0.4555), to_apply=%called_computation.25 }, execution_thread="sparsecore" %region_96.114 (scatter-add.48: s32[], scatter-add.49: s32[]) -> s32[] { %scatter-add.48 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.49 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1396 = s32[]{:T(128)S(7)} add(%scatter-add.48, %scatter-add.49), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1406 = s32[]{:T(128)S(7)} add(%scatter-add.48, %scatter-add.49), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.29.clone.clone (param_0.4555: s32[263], param_1.5343: s32[256], param_2.4512: s32[256]) -> s32[263] { - %param_0.4555 = s32[263]{0:T(512)} parameter(0) - %param_1.5343 = s32[256]{0:T(256)} parameter(1) - %reshape.3935 = s32[256]{0:T(256)} reshape(%param_1.5343), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %transpose.1112 = s32[256]{0:T(256)} transpose(%reshape.3935), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %param_2.4512 = s32[256]{0:T(256)} parameter(2) - %reshape.3936 = s32[256]{0:T(256)} reshape(%param_2.4512), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1113 = s32[256]{0:T(256)} transpose(%reshape.3936), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.243 = s32[263]{0:T(512)} scatter(%param_0.4555, %transpose.1112, %transpose.1113), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_96.114, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.26 (param_0.4556: s32[263], param_1.5344: s32[256], param_2.4513: s32[256]) -> s32[263] { +%fused_computation.29.clone.clone (param_0.4556: s32[263], param_1.5339: s32[256], param_2.4510: s32[256]) -> s32[263] { %param_0.4556 = s32[263]{0:T(512)} parameter(0) - %param_1.5344 = s32[256]{0:T(256)} parameter(1) - %param_2.4513 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.51 = s32[263]{0:T(512)} fusion(%param_0.4556, %param_1.5344, %param_2.4513), kind=kCustom, calls=%fused_computation.29.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5339 = s32[256]{0:T(256)} parameter(1) + %reshape.3915 = s32[256]{0:T(256)} reshape(%param_1.5339), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %transpose.1112 = s32[256]{0:T(256)} transpose(%reshape.3915), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %param_2.4510 = s32[256]{0:T(256)} parameter(2) + %reshape.3916 = s32[256]{0:T(256)} reshape(%param_2.4510), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1113 = s32[256]{0:T(256)} transpose(%reshape.3916), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.237 = s32[263]{0:T(512)} scatter(%param_0.4556, %transpose.1112, %transpose.1113), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_96.114, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.26 (param_0.4557: s32[263], param_1.5345: s32[256], param_2.4514: s32[256]) -> s32[263] { +%called_computation.26 (param_0.4557: s32[263], param_1.5340: s32[256], param_2.4511: s32[256]) -> s32[263] { %param_0.4557 = s32[263]{0:T(512)} parameter(0) - %param_1.5345 = s32[256]{0:T(256)} parameter(1) - %param_2.4514 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.52.cloned.1 = s32[263]{0:T(512)} call(%param_0.4557, %param_1.5345, %param_2.4514), to_apply=%called_computation.26, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + %param_1.5340 = s32[256]{0:T(256)} parameter(1) + %param_2.4511 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.51 = s32[263]{0:T(512)} fusion(%param_0.4557, %param_1.5340, %param_2.4511), kind=kCustom, calls=%fused_computation.29.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.6 (param_0.102: s32[263], param_1.154: s32[256], param_2.98: s32[256], param_3.3093: token[]) -> s32[263] { - %param_3.3093 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.26 (param_0.4558: s32[263], param_1.5341: s32[256], param_2.4512: s32[256]) -> s32[263] { + %param_0.4558 = s32[263]{0:T(512)} parameter(0) + %param_1.5341 = s32[256]{0:T(256)} parameter(1) + %param_2.4512 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.52.cloned.1 = s32[263]{0:T(512)} call(%param_0.4558, %param_1.5341, %param_2.4512), to_apply=%called_computation.26, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.6 (param_0.102: s32[263], param_1.154: s32[256], param_2.98: s32[256], param_3.3091: token[]) -> s32[263] { + %param_3.3091 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.102 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.154 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.98 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -950,47 +950,47 @@ StackFrames ROOT %scatter_offload_custom_fusion.52.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.52.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.6 (param_0.103: s32[263], param_1.155: s32[256], param_2.99: s32[256], param_3.3092: token[]) -> s32[263] { - %param_3.3092 = token[] parameter(3) +%async_computation.6 (param_0.103: s32[263], param_1.155: s32[256], param_2.99: s32[256], param_3.3090: token[]) -> s32[263] { + %param_3.3090 = token[] parameter(3) %param_0.103 = s32[263]{0:T(512)} parameter(0) %param_1.155 = s32[256]{0:T(256)} parameter(1) %param_2.99 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.20.cloned.1 = s32[263]{0:T(512)} call(%param_0.103, %param_1.155, %param_2.99, %param_3.3092), to_apply=%called_computation.6, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.20.cloned.1 = s32[263]{0:T(512)} call(%param_0.103, %param_1.155, %param_2.99, %param_3.3090), to_apply=%called_computation.6, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_102.120 (scatter-add.52: f32[], scatter-add.53: f32[]) -> f32[] { %scatter-add.52 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.53 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1399 = f32[]{:T(128)S(7)} add(%scatter-add.52, %scatter-add.53), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1409 = f32[]{:T(128)S(7)} add(%scatter-add.52, %scatter-add.53), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.30.clone.clone (param_0.4560: f32[9], param_1.5346: s32[256], param_2.4515: f32[256]) -> f32[9] { - %param_0.4560 = f32[9]{0:T(128)} parameter(0) - %param_1.5346 = s32[256]{0:T(256)} parameter(1) - %reshape.3937 = s32[256]{0:T(256)} reshape(%param_1.5346), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1114 = s32[256]{0:T(256)} transpose(%reshape.3937), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4515 = f32[256]{0:T(256)} parameter(2) - %reshape.3938 = f32[256]{0:T(256)} reshape(%param_2.4515), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1115 = f32[256]{0:T(256)} transpose(%reshape.3938), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.244 = f32[9]{0:T(128)} scatter(%param_0.4560, %transpose.1114, %transpose.1115), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_102.120, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.28 (param_0.4561: f32[9], param_1.5347: s32[256], param_2.4516: f32[256]) -> f32[9] { +%fused_computation.30.clone.clone (param_0.4561: f32[9], param_1.5342: s32[256], param_2.4513: f32[256]) -> f32[9] { %param_0.4561 = f32[9]{0:T(128)} parameter(0) - %param_1.5347 = s32[256]{0:T(256)} parameter(1) - %param_2.4516 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.53 = f32[9]{0:T(128)} fusion(%param_0.4561, %param_1.5347, %param_2.4516), kind=kCustom, calls=%fused_computation.30.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5342 = s32[256]{0:T(256)} parameter(1) + %reshape.3917 = s32[256]{0:T(256)} reshape(%param_1.5342), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1114 = s32[256]{0:T(256)} transpose(%reshape.3917), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4513 = f32[256]{0:T(256)} parameter(2) + %reshape.3918 = f32[256]{0:T(256)} reshape(%param_2.4513), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1115 = f32[256]{0:T(256)} transpose(%reshape.3918), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.238 = f32[9]{0:T(128)} scatter(%param_0.4561, %transpose.1114, %transpose.1115), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_102.120, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.28 (param_0.4562: f32[9], param_1.5348: s32[256], param_2.4517: f32[256]) -> f32[9] { +%called_computation.28 (param_0.4562: f32[9], param_1.5343: s32[256], param_2.4514: f32[256]) -> f32[9] { %param_0.4562 = f32[9]{0:T(128)} parameter(0) - %param_1.5348 = s32[256]{0:T(256)} parameter(1) - %param_2.4517 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.54.cloned.1 = f32[9]{0:T(128)} call(%param_0.4562, %param_1.5348, %param_2.4517), to_apply=%called_computation.28, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + %param_1.5343 = s32[256]{0:T(256)} parameter(1) + %param_2.4514 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.53 = f32[9]{0:T(128)} fusion(%param_0.4562, %param_1.5343, %param_2.4514), kind=kCustom, calls=%fused_computation.30.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.7 (param_0.105: f32[9], param_1.157: s32[256], param_2.101: f32[256], param_3.3101: token[]) -> f32[9] { - %param_3.3101 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.28 (param_0.4563: f32[9], param_1.5344: s32[256], param_2.4515: f32[256]) -> f32[9] { + %param_0.4563 = f32[9]{0:T(128)} parameter(0) + %param_1.5344 = s32[256]{0:T(256)} parameter(1) + %param_2.4515 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.54.cloned.1 = f32[9]{0:T(128)} call(%param_0.4563, %param_1.5344, %param_2.4515), to_apply=%called_computation.28, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.7 (param_0.105: f32[9], param_1.157: s32[256], param_2.101: f32[256], param_3.3099: token[]) -> f32[9] { + %param_3.3099 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.105 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.157 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.101 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -998,47 +998,47 @@ StackFrames ROOT %scatter_offload_custom_fusion.54.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.54.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.7 (param_0.106: f32[9], param_1.158: s32[256], param_2.102: f32[256], param_3.3100: token[]) -> f32[9] { - %param_3.3100 = token[] parameter(3) +%async_computation.7 (param_0.106: f32[9], param_1.158: s32[256], param_2.102: f32[256], param_3.3098: token[]) -> f32[9] { + %param_3.3098 = token[] parameter(3) %param_0.106 = f32[9]{0:T(128)} parameter(0) %param_1.158 = s32[256]{0:T(256)} parameter(1) %param_2.102 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.23.cloned.1 = f32[9]{0:T(128)} call(%param_0.106, %param_1.158, %param_2.102, %param_3.3100), to_apply=%called_computation.7, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.23.cloned.1 = f32[9]{0:T(128)} call(%param_0.106, %param_1.158, %param_2.102, %param_3.3098), to_apply=%called_computation.7, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_104.122 (scatter-add.83: s32[], scatter-add.84: s32[]) -> s32[] { %scatter-add.83 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.84 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1400 = s32[]{:T(128)S(7)} add(%scatter-add.83, %scatter-add.84), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1410 = s32[]{:T(128)S(7)} add(%scatter-add.83, %scatter-add.84), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.31.clone.clone (param_0.4565: s32[263], param_1.5349: s32[8], param_2.4518: s32[8]) -> s32[263] { - %param_0.4565 = s32[263]{0:T(512)} parameter(0) - %param_1.5349 = s32[8]{0:T(128)} parameter(1) - %reshape.3939 = s32[8]{0:T(128)} reshape(%param_1.5349), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %transpose.1116 = s32[8]{0:T(128)} transpose(%reshape.3939), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %param_2.4518 = s32[8]{0:T(128)} parameter(2) - %reshape.3940 = s32[8]{0:T(128)} reshape(%param_2.4518), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1117 = s32[8]{0:T(128)} transpose(%reshape.3940), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.245 = s32[263]{0:T(512)} scatter(%param_0.4565, %transpose.1116, %transpose.1117), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_104.122, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.30 (param_0.4566: s32[263], param_1.5350: s32[8], param_2.4519: s32[8]) -> s32[263] { +%fused_computation.31.clone.clone (param_0.4566: s32[263], param_1.5345: s32[8], param_2.4516: s32[8]) -> s32[263] { %param_0.4566 = s32[263]{0:T(512)} parameter(0) - %param_1.5350 = s32[8]{0:T(128)} parameter(1) - %param_2.4519 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.55 = s32[263]{0:T(512)} fusion(%param_0.4566, %param_1.5350, %param_2.4519), kind=kCustom, calls=%fused_computation.31.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5345 = s32[8]{0:T(128)} parameter(1) + %reshape.3919 = s32[8]{0:T(128)} reshape(%param_1.5345), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %transpose.1116 = s32[8]{0:T(128)} transpose(%reshape.3919), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %param_2.4516 = s32[8]{0:T(128)} parameter(2) + %reshape.3920 = s32[8]{0:T(128)} reshape(%param_2.4516), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + %transpose.1117 = s32[8]{0:T(128)} transpose(%reshape.3920), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} + ROOT %scatter-add.239 = s32[263]{0:T(512)} scatter(%param_0.4566, %transpose.1116, %transpose.1117), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_104.122, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.30 (param_0.4567: s32[263], param_1.5351: s32[8], param_2.4520: s32[8]) -> s32[263] { +%called_computation.30 (param_0.4567: s32[263], param_1.5346: s32[8], param_2.4517: s32[8]) -> s32[263] { %param_0.4567 = s32[263]{0:T(512)} parameter(0) - %param_1.5351 = s32[8]{0:T(128)} parameter(1) - %param_2.4520 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.56.cloned.1 = s32[263]{0:T(512)} call(%param_0.4567, %param_1.5351, %param_2.4520), to_apply=%called_computation.30, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + %param_1.5346 = s32[8]{0:T(128)} parameter(1) + %param_2.4517 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.55 = s32[263]{0:T(512)} fusion(%param_0.4567, %param_1.5346, %param_2.4517), kind=kCustom, calls=%fused_computation.31.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +}, execution_thread="sparsecore" + +%async_computation.30 (param_0.4568: s32[263], param_1.5347: s32[8], param_2.4518: s32[8]) -> s32[263] { + %param_0.4568 = s32[263]{0:T(512)} parameter(0) + %param_1.5347 = s32[8]{0:T(128)} parameter(1) + %param_2.4518 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.56.cloned.1 = s32[263]{0:T(512)} call(%param_0.4568, %param_1.5347, %param_2.4518), to_apply=%called_computation.30, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.8 (param_0.108: s32[263], param_1.160: s32[8], param_2.104: s32[8], param_3.3109: token[]) -> s32[263] { - %param_3.3109 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%called_computation.8 (param_0.108: s32[263], param_1.160: s32[8], param_2.104: s32[8], param_3.3107: token[]) -> s32[263] { + %param_3.3107 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.108 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.160 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.104 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1046,47 +1046,47 @@ StackFrames ROOT %scatter_offload_custom_fusion.56.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.56.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.8 (param_0.109: s32[263], param_1.161: s32[8], param_2.105: s32[8], param_3.3108: token[]) -> s32[263] { - %param_3.3108 = token[] parameter(3) +%async_computation.8 (param_0.109: s32[263], param_1.161: s32[8], param_2.105: s32[8], param_3.3106: token[]) -> s32[263] { + %param_3.3106 = token[] parameter(3) %param_0.109 = s32[263]{0:T(512)} parameter(0) %param_1.161 = s32[8]{0:T(128)} parameter(1) %param_2.105 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.26.cloned.1 = s32[263]{0:T(512)} call(%param_0.109, %param_1.161, %param_2.105, %param_3.3108), to_apply=%called_computation.8, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.26.cloned.1 = s32[263]{0:T(512)} call(%param_0.109, %param_1.161, %param_2.105, %param_3.3106), to_apply=%called_computation.8, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_14.20 (scatter-add.0: s32[], scatter-add.1: s32[]) -> s32[] { %scatter-add.0 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.1 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1312 = s32[]{:T(128)S(7)} add(%scatter-add.0, %scatter-add.1), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.1322 = s32[]{:T(128)S(7)} add(%scatter-add.0, %scatter-add.1), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.17.clone.clone.clone (param_0.4570: s32[256], param_1.5352: s32[4096], param_2.4521: s32[4096]) -> s32[256] { - %param_0.4570 = s32[256]{0:T(256)} parameter(0) - %param_1.5352 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3941 = s32[4096]{0:T(1024)} reshape(%param_1.5352), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} - %transpose.1118 = s32[4096]{0:T(1024)} transpose(%reshape.3941), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} - %param_2.4521 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3942 = s32[4096]{0:T(1024)} reshape(%param_2.4521), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - %transpose.1119 = s32[4096]{0:T(1024)} transpose(%reshape.3942), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.246 = s32[256]{0:T(256)} scatter(%param_0.4570, %transpose.1118, %transpose.1119), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_14.20, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.32 (param_0.4571: s32[256], param_1.5353: s32[4096], param_2.4522: s32[4096]) -> s32[256] { +%fused_computation.17.clone.clone.clone (param_0.4571: s32[256], param_1.5348: s32[4096], param_2.4519: s32[4096]) -> s32[256] { %param_0.4571 = s32[256]{0:T(256)} parameter(0) - %param_1.5353 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4522 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.57 = s32[256]{0:T(256)} fusion(%param_0.4571, %param_1.5353, %param_2.4522), kind=kCustom, calls=%fused_computation.17.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5348 = s32[4096]{0:T(1024)} parameter(1) + %reshape.3921 = s32[4096]{0:T(1024)} reshape(%param_1.5348), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} + %transpose.1118 = s32[4096]{0:T(1024)} transpose(%reshape.3921), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} + %param_2.4519 = s32[4096]{0:T(1024)} parameter(2) + %reshape.3922 = s32[4096]{0:T(1024)} reshape(%param_2.4519), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + %transpose.1119 = s32[4096]{0:T(1024)} transpose(%reshape.3922), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.240 = s32[256]{0:T(256)} scatter(%param_0.4571, %transpose.1118, %transpose.1119), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_14.20, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.32 (param_0.4572: s32[256], param_1.5354: s32[4096], param_2.4523: s32[4096]) -> s32[256] { +%called_computation.32 (param_0.4572: s32[256], param_1.5349: s32[4096], param_2.4520: s32[4096]) -> s32[256] { %param_0.4572 = s32[256]{0:T(256)} parameter(0) - %param_1.5354 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4523 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.58.cloned.1 = s32[256]{0:T(256)} call(%param_0.4572, %param_1.5354, %param_2.4523), to_apply=%called_computation.32, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} + %param_1.5349 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4520 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.57 = s32[256]{0:T(256)} fusion(%param_0.4572, %param_1.5349, %param_2.4520), kind=kCustom, calls=%fused_computation.17.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.9 (param_0.111: s32[256], param_1.163: s32[4096], param_2.107: s32[4096], param_3.3087: token[]) -> s32[256] { - %param_3.3087 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.32 (param_0.4573: s32[256], param_1.5350: s32[4096], param_2.4521: s32[4096]) -> s32[256] { + %param_0.4573 = s32[256]{0:T(256)} parameter(0) + %param_1.5350 = s32[4096]{0:T(1024)} parameter(1) + %param_2.4521 = s32[4096]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.58.cloned.1 = s32[256]{0:T(256)} call(%param_0.4573, %param_1.5350, %param_2.4521), to_apply=%called_computation.32, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.9 (param_0.111: s32[256], param_1.163: s32[4096], param_2.107: s32[4096], param_3.3085: token[]) -> s32[256] { + %param_3.3085 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.111 = s32[256]{0:T(256)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.163 = s32[4096]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.107 = s32[4096]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1094,57 +1094,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.58.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.58.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.9 (param_0.112: s32[256], param_1.164: s32[4096], param_2.108: s32[4096], param_3.3086: token[]) -> s32[256] { - %param_3.3086 = token[] parameter(3) +%async_computation.9 (param_0.112: s32[256], param_1.164: s32[4096], param_2.108: s32[4096], param_3.3084: token[]) -> s32[256] { + %param_3.3084 = token[] parameter(3) %param_0.112 = s32[256]{0:T(256)} parameter(0) %param_1.164 = s32[4096]{0:T(1024)} parameter(1) %param_2.108 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.29.cloned.1 = s32[256]{0:T(256)} call(%param_0.112, %param_1.164, %param_2.108, %param_3.3086), to_apply=%called_computation.9, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.29.cloned.1 = s32[256]{0:T(256)} call(%param_0.112, %param_1.164, %param_2.108, %param_3.3084), to_apply=%called_computation.9, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.33 (param_0.4573: s32[263]) -> s32[263] { - %param_0.4573 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2093 = s32[263]{0:T(512)} copy(%param_0.4573), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.33 (param_0.4574: s32[263]) -> s32[263] { + %param_0.4574 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2093 = s32[263]{0:T(512)} copy(%param_0.4574), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.33 (param_0.4574: s32[263]) -> s32[263] { - %param_0.4574 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2094.cloned.1 = s32[263]{0:T(512)} call(%param_0.4574), to_apply=%called_computation.33 +%async_computation.33 (param_0.4575: s32[263]) -> s32[263] { + %param_0.4575 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2094.cloned.1 = s32[263]{0:T(512)} call(%param_0.4575), to_apply=%called_computation.33 }, execution_thread="sparsecore" %region_20.26.clone.1 (scatter-add.141: s32[], scatter-add.142: s32[]) -> s32[] { %scatter-add.141 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.142 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2463 = s32[]{:T(128)S(7)} add(%scatter-add.141, %scatter-add.142), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.2471 = s32[]{:T(128)S(7)} add(%scatter-add.141, %scatter-add.142), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.18.clone.clone.clone (param_0.4575: s32[263], param_1.5355: s32[256], param_2.4524: s32[256]) -> s32[263] { - %param_0.4575 = s32[263]{0:T(512)} parameter(0) - %param_1.5355 = s32[256]{0:T(256)} parameter(1) - %reshape.3943 = s32[256]{0:T(256)} reshape(%param_1.5355), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1120 = s32[256]{0:T(256)} transpose(%reshape.3943), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4524 = s32[256]{0:T(256)} parameter(2) - %reshape.3944 = s32[256]{0:T(256)} reshape(%param_2.4524) - %transpose.1121 = s32[256]{0:T(256)} transpose(%reshape.3944), dimensions={0} - ROOT %scatter-add.247 = s32[263]{0:T(512)} scatter(%param_0.4575, %transpose.1120, %transpose.1121), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_20.26.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.34 (param_0.4576: s32[263], param_1.5356: s32[256], param_2.4525: s32[256]) -> s32[263] { +%fused_computation.18.clone.clone.clone (param_0.4576: s32[263], param_1.5351: s32[256], param_2.4522: s32[256]) -> s32[263] { %param_0.4576 = s32[263]{0:T(512)} parameter(0) - %param_1.5356 = s32[256]{0:T(256)} parameter(1) - %param_2.4525 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.59 = s32[263]{0:T(512)} fusion(%param_0.4576, %param_1.5356, %param_2.4525), kind=kCustom, calls=%fused_computation.18.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5351 = s32[256]{0:T(256)} parameter(1) + %reshape.3923 = s32[256]{0:T(256)} reshape(%param_1.5351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1120 = s32[256]{0:T(256)} transpose(%reshape.3923), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4522 = s32[256]{0:T(256)} parameter(2) + %reshape.3924 = s32[256]{0:T(256)} reshape(%param_2.4522) + %transpose.1121 = s32[256]{0:T(256)} transpose(%reshape.3924), dimensions={0} + ROOT %scatter-add.241 = s32[263]{0:T(512)} scatter(%param_0.4576, %transpose.1120, %transpose.1121), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_20.26.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.34 (param_0.4577: s32[263], param_1.5357: s32[256], param_2.4526: s32[256]) -> s32[263] { +%called_computation.34 (param_0.4577: s32[263], param_1.5352: s32[256], param_2.4523: s32[256]) -> s32[263] { %param_0.4577 = s32[263]{0:T(512)} parameter(0) - %param_1.5357 = s32[256]{0:T(256)} parameter(1) - %param_2.4526 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.60.cloned.1 = s32[263]{0:T(512)} call(%param_0.4577, %param_1.5357, %param_2.4526), to_apply=%called_computation.34, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5352 = s32[256]{0:T(256)} parameter(1) + %param_2.4523 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.59 = s32[263]{0:T(512)} fusion(%param_0.4577, %param_1.5352, %param_2.4523), kind=kCustom, calls=%fused_computation.18.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.10 (param_0.114: s32[263], param_1.166: s32[256], param_2.110: s32[256], param_3.3089: token[]) -> s32[263] { - %param_3.3089 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.34 (param_0.4578: s32[263], param_1.5353: s32[256], param_2.4524: s32[256]) -> s32[263] { + %param_0.4578 = s32[263]{0:T(512)} parameter(0) + %param_1.5353 = s32[256]{0:T(256)} parameter(1) + %param_2.4524 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.60.cloned.1 = s32[263]{0:T(512)} call(%param_0.4578, %param_1.5353, %param_2.4524), to_apply=%called_computation.34, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.10 (param_0.114: s32[263], param_1.166: s32[256], param_2.110: s32[256], param_3.3087: token[]) -> s32[263] { + %param_3.3087 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.114 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.166 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.110 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1154,57 +1154,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.60.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.60.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.10 (param_0.115: s32[263], param_1.167: s32[256], param_2.111: s32[256], param_3.3088: token[]) -> s32[263] { - %param_3.3088 = token[] parameter(3) +%async_computation.10 (param_0.115: s32[263], param_1.167: s32[256], param_2.111: s32[256], param_3.3086: token[]) -> s32[263] { + %param_3.3086 = token[] parameter(3) %param_0.115 = s32[263]{0:T(512)} parameter(0) %param_1.167 = s32[256]{0:T(256)} parameter(1) %param_2.111 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.32.cloned.1 = s32[263]{0:T(512)} call(%param_0.115, %param_1.167, %param_2.111, %param_3.3088), to_apply=%called_computation.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.32.cloned.1 = s32[263]{0:T(512)} call(%param_0.115, %param_1.167, %param_2.111, %param_3.3086), to_apply=%called_computation.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.35 (param_0.4578: f32[9]) -> f32[9] { - %param_0.4578 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2095 = f32[9]{0:T(128)} copy(%param_0.4578), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.35 (param_0.4579: f32[9]) -> f32[9] { + %param_0.4579 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2095 = f32[9]{0:T(128)} copy(%param_0.4579), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.35 (param_0.4579: f32[9]) -> f32[9] { - %param_0.4579 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2096.cloned.1 = f32[9]{0:T(128)} call(%param_0.4579), to_apply=%called_computation.35 +%async_computation.35 (param_0.4580: f32[9]) -> f32[9] { + %param_0.4580 = f32[9]{0:T(128)} parameter(0) + ROOT %copy.2096.cloned.1 = f32[9]{0:T(128)} call(%param_0.4580), to_apply=%called_computation.35 }, execution_thread="sparsecore" %region_26.33.clone.1 (scatter-add.145: f32[], scatter-add.146: f32[]) -> f32[] { %scatter-add.145 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.146 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2465 = f32[]{:T(128)S(7)} add(%scatter-add.145, %scatter-add.146), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.19.clone.clone.clone (param_0.4580: f32[9], param_1.5358: s32[256], param_2.4527: f32[256]) -> f32[9] { - %param_0.4580 = f32[9]{0:T(128)} parameter(0) - %param_1.5358 = s32[256]{0:T(256)} parameter(1) - %reshape.3945 = s32[256]{0:T(256)} reshape(%param_1.5358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1122 = s32[256]{0:T(256)} transpose(%reshape.3945), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4527 = f32[256]{0:T(256)} parameter(2) - %reshape.3946 = f32[256]{0:T(256)} reshape(%param_2.4527) - %transpose.1123 = f32[256]{0:T(256)} transpose(%reshape.3946), dimensions={0} - ROOT %scatter-add.248 = f32[9]{0:T(128)} scatter(%param_0.4580, %transpose.1122, %transpose.1123), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_26.33.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %add.2473 = f32[]{:T(128)S(7)} add(%scatter-add.145, %scatter-add.146), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.36 (param_0.4581: f32[9], param_1.5359: s32[256], param_2.4528: f32[256]) -> f32[9] { +%fused_computation.19.clone.clone.clone (param_0.4581: f32[9], param_1.5354: s32[256], param_2.4525: f32[256]) -> f32[9] { %param_0.4581 = f32[9]{0:T(128)} parameter(0) - %param_1.5359 = s32[256]{0:T(256)} parameter(1) - %param_2.4528 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.61 = f32[9]{0:T(128)} fusion(%param_0.4581, %param_1.5359, %param_2.4528), kind=kCustom, calls=%fused_computation.19.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5354 = s32[256]{0:T(256)} parameter(1) + %reshape.3925 = s32[256]{0:T(256)} reshape(%param_1.5354), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1122 = s32[256]{0:T(256)} transpose(%reshape.3925), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4525 = f32[256]{0:T(256)} parameter(2) + %reshape.3926 = f32[256]{0:T(256)} reshape(%param_2.4525) + %transpose.1123 = f32[256]{0:T(256)} transpose(%reshape.3926), dimensions={0} + ROOT %scatter-add.242 = f32[9]{0:T(128)} scatter(%param_0.4581, %transpose.1122, %transpose.1123), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_26.33.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.36 (param_0.4582: f32[9], param_1.5360: s32[256], param_2.4529: f32[256]) -> f32[9] { +%called_computation.36 (param_0.4582: f32[9], param_1.5355: s32[256], param_2.4526: f32[256]) -> f32[9] { %param_0.4582 = f32[9]{0:T(128)} parameter(0) - %param_1.5360 = s32[256]{0:T(256)} parameter(1) - %param_2.4529 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.62.cloned.1 = f32[9]{0:T(128)} call(%param_0.4582, %param_1.5360, %param_2.4529), to_apply=%called_computation.36, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5355 = s32[256]{0:T(256)} parameter(1) + %param_2.4526 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.61 = f32[9]{0:T(128)} fusion(%param_0.4582, %param_1.5355, %param_2.4526), kind=kCustom, calls=%fused_computation.19.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.11 (param_0.117: f32[9], param_1.169: s32[256], param_2.113: f32[256], param_3.3095: token[]) -> f32[9] { - %param_3.3095 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.36 (param_0.4583: f32[9], param_1.5356: s32[256], param_2.4527: f32[256]) -> f32[9] { + %param_0.4583 = f32[9]{0:T(128)} parameter(0) + %param_1.5356 = s32[256]{0:T(256)} parameter(1) + %param_2.4527 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.62.cloned.1 = f32[9]{0:T(128)} call(%param_0.4583, %param_1.5356, %param_2.4527), to_apply=%called_computation.36, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.11 (param_0.117: f32[9], param_1.169: s32[256], param_2.113: f32[256], param_3.3093: token[]) -> f32[9] { + %param_3.3093 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.117 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.169 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.113 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1214,57 +1214,57 @@ StackFrames ROOT %scatter_offload_custom_fusion.62.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.62.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.11 (param_0.118: f32[9], param_1.170: s32[256], param_2.114: f32[256], param_3.3094: token[]) -> f32[9] { - %param_3.3094 = token[] parameter(3) +%async_computation.11 (param_0.118: f32[9], param_1.170: s32[256], param_2.114: f32[256], param_3.3092: token[]) -> f32[9] { + %param_3.3092 = token[] parameter(3) %param_0.118 = f32[9]{0:T(128)} parameter(0) %param_1.170 = s32[256]{0:T(256)} parameter(1) %param_2.114 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.35.cloned.1 = f32[9]{0:T(128)} call(%param_0.118, %param_1.170, %param_2.114, %param_3.3094), to_apply=%called_computation.11, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.35.cloned.1 = f32[9]{0:T(128)} call(%param_0.118, %param_1.170, %param_2.114, %param_3.3092), to_apply=%called_computation.11, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.37 (param_0.4583: s32[263]) -> s32[263] { - %param_0.4583 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2097 = s32[263]{0:T(512)} copy(%param_0.4583), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} +%called_computation.37 (param_0.4584: s32[263]) -> s32[263] { + %param_0.4584 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2097 = s32[263]{0:T(512)} copy(%param_0.4584), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%async_computation.37 (param_0.4584: s32[263]) -> s32[263] { - %param_0.4584 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2098.cloned.1 = s32[263]{0:T(512)} call(%param_0.4584), to_apply=%called_computation.37 +%async_computation.37 (param_0.4585: s32[263]) -> s32[263] { + %param_0.4585 = s32[263]{0:T(512)} parameter(0) + ROOT %copy.2098.cloned.1 = s32[263]{0:T(512)} call(%param_0.4585), to_apply=%called_computation.37 }, execution_thread="sparsecore" %region_28.35.clone.1 (scatter-add.149: s32[], scatter-add.150: s32[]) -> s32[] { %scatter-add.149 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} %scatter-add.150 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2467 = s32[]{:T(128)S(7)} add(%scatter-add.149, %scatter-add.150), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + ROOT %add.2475 = s32[]{:T(128)S(7)} add(%scatter-add.149, %scatter-add.150), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%fused_computation.20.clone.clone.clone (param_0.4585: s32[263], param_1.5361: s32[8], param_2.4530: s32[8]) -> s32[263] { - %param_0.4585 = s32[263]{0:T(512)} parameter(0) - %param_1.5361 = s32[8]{0:T(128)} parameter(1) - %reshape.3947 = s32[8]{0:T(128)} reshape(%param_1.5361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1124 = s32[8]{0:T(128)} transpose(%reshape.3947), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4530 = s32[8]{0:T(128)} parameter(2) - %reshape.3948 = s32[8]{0:T(128)} reshape(%param_2.4530) - %transpose.1125 = s32[8]{0:T(128)} transpose(%reshape.3948), dimensions={0} - ROOT %scatter-add.249 = s32[263]{0:T(512)} scatter(%param_0.4585, %transpose.1124, %transpose.1125), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_28.35.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.38 (param_0.4586: s32[263], param_1.5362: s32[8], param_2.4531: s32[8]) -> s32[263] { +%fused_computation.20.clone.clone.clone (param_0.4586: s32[263], param_1.5357: s32[8], param_2.4528: s32[8]) -> s32[263] { %param_0.4586 = s32[263]{0:T(512)} parameter(0) - %param_1.5362 = s32[8]{0:T(128)} parameter(1) - %param_2.4531 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.63 = s32[263]{0:T(512)} fusion(%param_0.4586, %param_1.5362, %param_2.4531), kind=kCustom, calls=%fused_computation.20.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} + %param_1.5357 = s32[8]{0:T(128)} parameter(1) + %reshape.3927 = s32[8]{0:T(128)} reshape(%param_1.5357), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1124 = s32[8]{0:T(128)} transpose(%reshape.3927), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4528 = s32[8]{0:T(128)} parameter(2) + %reshape.3928 = s32[8]{0:T(128)} reshape(%param_2.4528) + %transpose.1125 = s32[8]{0:T(128)} transpose(%reshape.3928), dimensions={0} + ROOT %scatter-add.243 = s32[263]{0:T(512)} scatter(%param_0.4586, %transpose.1124, %transpose.1125), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_28.35.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.38 (param_0.4587: s32[263], param_1.5363: s32[8], param_2.4532: s32[8]) -> s32[263] { +%called_computation.38 (param_0.4587: s32[263], param_1.5358: s32[8], param_2.4529: s32[8]) -> s32[263] { %param_0.4587 = s32[263]{0:T(512)} parameter(0) - %param_1.5363 = s32[8]{0:T(128)} parameter(1) - %param_2.4532 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.64.cloned.1 = s32[263]{0:T(512)} call(%param_0.4587, %param_1.5363, %param_2.4532), to_apply=%called_computation.38, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + %param_1.5358 = s32[8]{0:T(128)} parameter(1) + %param_2.4529 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.63 = s32[263]{0:T(512)} fusion(%param_0.4587, %param_1.5358, %param_2.4529), kind=kCustom, calls=%fused_computation.20.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} }, execution_thread="sparsecore" -%called_computation.12 (param_0.120: s32[263], param_1.172: s32[8], param_2.116: s32[8], param_3.3103: token[]) -> s32[263] { - %param_3.3103 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} +%async_computation.38 (param_0.4588: s32[263], param_1.5359: s32[8], param_2.4530: s32[8]) -> s32[263] { + %param_0.4588 = s32[263]{0:T(512)} parameter(0) + %param_1.5359 = s32[8]{0:T(128)} parameter(1) + %param_2.4530 = s32[8]{0:T(128)} parameter(2) + ROOT %scatter_offload_custom_fusion.64.cloned.1 = s32[263]{0:T(512)} call(%param_0.4588, %param_1.5359, %param_2.4530), to_apply=%called_computation.38, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.12 (param_0.120: s32[263], param_1.172: s32[8], param_2.116: s32[8], param_3.3101: token[]) -> s32[263] { + %param_3.3101 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} %param_0.120 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_1.172 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} %param_2.116 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} @@ -1274,12 +1274,12 @@ StackFrames ROOT %scatter_offload_custom_fusion.64.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.64.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.12 (param_0.121: s32[263], param_1.173: s32[8], param_2.117: s32[8], param_3.3102: token[]) -> s32[263] { - %param_3.3102 = token[] parameter(3) +%async_computation.12 (param_0.121: s32[263], param_1.173: s32[8], param_2.117: s32[8], param_3.3100: token[]) -> s32[263] { + %param_3.3100 = token[] parameter(3) %param_0.121 = s32[263]{0:T(512)} parameter(0) %param_1.173 = s32[8]{0:T(128)} parameter(1) %param_2.117 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.38.cloned.1 = s32[263]{0:T(512)} call(%param_0.121, %param_1.173, %param_2.117, %param_3.3102), to_apply=%called_computation.12, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.38.cloned.1 = s32[263]{0:T(512)} call(%param_0.121, %param_1.173, %param_2.117, %param_3.3100), to_apply=%called_computation.12, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" %region_154.179 (reduce_sum.431: f32[], reduce_sum.254: f32[]) -> f32[] { @@ -1288,18 +1288,18 @@ StackFrames ROOT %reduce_sum.258 = f32[]{:T(128)} add(%reduce_sum.431, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.467 (param_0.4170: f32[3,1536,128,192]) -> f32[] { - %param_0.4170 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(0) - %bitcast.672 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4170), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.564 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.672, %bitcast.672), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5105 = f32[]{:T(128)} constant(0) - ROOT %reduce.669 = f32[]{:T(128)} reduce(%square.564, %constant.5105), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.466 (param_0.4171: f32[3,1536,128,192]) -> f32[] { + %param_0.4171 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(0) + %bitcast.670 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4171), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3786 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.670, %bitcast.670), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5085 = f32[]{:T(128)} constant(0) + ROOT %reduce.669 = f32[]{:T(128)} reduce(%mul.3786, %constant.5085), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.468 (param_0.1421: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { - %param_0.1421 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) - %copy.1550 = bf16[1536,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1421), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\']"} - ROOT %bitcast.673 = bf16[3,1536,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1550), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} +%fused_computation.467 (param_0.1419: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { + %param_0.1419 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %copy.1550 = bf16[1536,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1419), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\']"} + ROOT %bitcast.671 = bf16[3,1536,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1550), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} } %region_221.246 (reduce_sum.893: f32[], reduce_sum.603: f32[]) -> f32[] { @@ -1314,55 +1314,55 @@ StackFrames ROOT %reduce_sum.450 = f32[]{:T(128)} add(%reduce_sum.655, %reduce_sum.449), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.469 (param_0.4140: f32[1536,3,128,192], param_1.5025: f32[], param_2.4298: f32[], param_3.2951: f32[], param_4.2203: f32[1536,3,128,192], param_5.2006: f32[], param_6.1443: f32[3,1536,128,192], param_7.1124: pred[], param_8.889: f32[1536,3,128,192]) -> (f32[], f32[1536,3,128,192], f32[1536,3,128,192], f32[1536,3,128,192], f32[]) { - %param_0.4140 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) - %param_3.2951 = f32[]{:T(128)S(6)} parameter(3) - %mul.4727.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.468 (param_0.4141: f32[1536,3,128,192], param_1.5021: f32[], param_2.4296: f32[], param_3.2949: f32[], param_4.2203: f32[1536,3,128,192], param_5.2006: f32[], param_6.1443: f32[3,1536,128,192], param_7.1124: pred[], param_8.889: f32[1536,3,128,192]) -> (f32[], f32[1536,3,128,192], f32[1536,3,128,192], f32[1536,3,128,192], f32[]) { + %param_0.4141 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %param_3.2949 = f32[]{:T(128)S(6)} parameter(3) + %mul.5031.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2949), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1124 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2165.clone.1 = pred[1536,3,128,192]{2,3,1,0:T(8,128)(4,1)} broadcast(%param_7.1124), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2121.clone.1 = pred[1536,3,128,192]{2,3,1,0:T(8,128)(4,1)} broadcast(%param_7.1124), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1443 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(6) - %bitcast.1374.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_6.1443), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1372.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_6.1443), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} %param_5.2006 = f32[]{:T(128)} parameter(5) - %div.2575.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_5.2006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2574.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.1374.clone.1, %div.2575.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2164.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2165.clone.1, %bitcast.1374.clone.1, %div.2574.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4864.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4279.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4864.clone.1), dimensions={}, metadata={op_name="broadcast.334"} - %mul.4733.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2164.clone.1, %broadcast.4279.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2565.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_5.2006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2564.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.1372.clone.1, %div.2565.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2120.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2121.clone.1, %bitcast.1372.clone.1, %div.2564.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4844.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4252.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4844.clone.1), dimensions={}, metadata={op_name="broadcast.334"} + %mul.5037.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.889 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(8) - %constant.4868.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4734.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4868.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4732.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.4734.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3443.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4733.clone.1, %mul.4732.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4298 = f32[]{:T(128)S(6)} parameter(2) - %div.2571.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4298), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.399.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2164.clone.1, %select_n.2164.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4867.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4731.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4867.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4729.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.4731.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4848.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.5038.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4848.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5036.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.5038.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3429.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.5037.clone.1, %mul.5036.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4296 = f32[]{:T(128)S(6)} parameter(2) + %div.2561.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4296), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.399.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %select_n.2120.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4847.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.5035.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4847.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5033.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.5035.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2203 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(4) - %constant.4866.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4730.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4866.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4728.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.4730.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3442.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4729.clone.1, %mul.4728.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5025 = f32[]{:T(128)S(6)} parameter(1) - %div.2570.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5025), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2569.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3442.clone.1, %div.2570.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.157.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} sqrt(%div.2569.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4865.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3441.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4865.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3440.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.157.clone.1, %add.3441.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1293.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%div.2571.clone.1, %add.3440.clone.1), metadata={op_name="multiply.290"} - %div.2568.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3443.clone.1, %multiply.1293.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4726.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4140, %broadcast.4279.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3439.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2568.clone.1, %mul.4726.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4725.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.4727.clone.1, %add.3439.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3438.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4140, %mul.4725.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.565 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3438.clone.1, %add.3438.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5075 = f32[]{:T(128)} constant(0) - %reduce.670 = f32[]{:T(128)} reduce(%square.565, %constant.5075), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.671.clone.1 = f32[]{:T(128)} reduce(%integer_pow.399.clone.1, %constant.5075), dimensions={0,1,2,3}, to_apply=%region_187.212, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.660 = (f32[]{:T(128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.670, %add.3438.clone.1, %add.3442.clone.1, %add.3443.clone.1, %reduce.671.clone.1) + %constant.4846.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.5034.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4846.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5032.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.5034.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3428.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.5033.clone.1, %mul.5032.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5021 = f32[]{:T(128)S(6)} parameter(1) + %div.2560.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5021), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2559.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3428.clone.1, %div.2560.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.157.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} sqrt(%div.2559.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4845.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3427.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4845.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3426.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.157.clone.1, %add.3427.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1293.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%div.2561.clone.1, %add.3426.clone.1), metadata={op_name="multiply.290"} + %div.2558.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3429.clone.1, %multiply.1293.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.5030.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4141, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3425.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2558.clone.1, %mul.5030.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.5029.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.5031.clone.1, %add.3425.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3424.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4141, %mul.5029.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.330 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3424.clone.1, %add.3424.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5055 = f32[]{:T(128)} constant(0) + %reduce.670 = f32[]{:T(128)} reduce(%square.330, %constant.5055), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.671.clone.1 = f32[]{:T(128)} reduce(%integer_pow.399.clone.1, %constant.5055), dimensions={0,1,2,3}, to_apply=%region_187.212, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.656 = (f32[]{:T(128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.670, %add.3424.clone.1, %add.3428.clone.1, %add.3429.clone.1, %reduce.671.clone.1) } %region_160.185 (reduce_sum.473: f32[], reduce_sum.293: f32[]) -> f32[] { @@ -1377,19 +1377,19 @@ StackFrames ROOT %reduce_sum.461 = f32[]{:T(128)} add(%reduce_sum.459, %reduce_sum.460), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.495 (param_0.4166: bf16[256,512,512], param_1.5047: bf16[256,512,512]) -> (f32[], f32[]) { - %param_0.4166 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) - %broadcast_in_dim.1358 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4166), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.695 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.570 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.695, %bitcast.695), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5101 = f32[]{:T(128)} constant(0) - %reduce.672 = f32[]{:T(128)} reduce(%square.570, %constant.5101), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.5047 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(1) - %broadcast_in_dim.1366.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_1.5047), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.703.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1366.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.576.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.703.clone.1, %bitcast.703.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.674.clone.1 = f32[]{:T(128)} reduce(%square.576.clone.1, %constant.5101), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.767 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.672, %reduce.674.clone.1) +%fused_computation.494 (param_0.4167: bf16[256,512,512], param_1.5043: bf16[256,512,512]) -> (f32[], f32[]) { + %param_0.4167 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %broadcast_in_dim.1245 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4167), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.693 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3815 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.693, %bitcast.693), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5081 = f32[]{:T(128)} constant(0) + %reduce.672 = f32[]{:T(128)} reduce(%mul.3815, %constant.5081), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.5043 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(1) + %broadcast_in_dim.1253.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_1.5043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.701.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1253.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3821.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.701.clone.1, %bitcast.701.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %reduce.674.clone.1 = f32[]{:T(128)} reduce(%mul.3821.clone.1, %constant.5081), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.763 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.672, %reduce.674.clone.1) } %region_159.184 (reduce_sum.466: f32[], reduce_sum.279: f32[]) -> f32[] { @@ -1398,13 +1398,13 @@ StackFrames ROOT %reduce_sum.286 = f32[]{:T(128)} add(%reduce_sum.466, %reduce_sum.279), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.497 (param_0.4165: bf16[256,512,512]) -> f32[] { - %param_0.4165 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) - %broadcast_in_dim.1362 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4165), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.699 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.573 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.699, %bitcast.699), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5100 = f32[]{:T(128)} constant(0) - ROOT %reduce.673 = f32[]{:T(128)} reduce(%square.573, %constant.5100), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.496 (param_0.4166: bf16[256,512,512]) -> f32[] { + %param_0.4166 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %broadcast_in_dim.1249 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4166), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.697 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3818 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.697, %bitcast.697), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5080 = f32[]{:T(128)} constant(0) + ROOT %reduce.673 = f32[]{:T(128)} reduce(%mul.3818, %constant.5080), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_227.252 (reduce_sum.935: f32[], reduce_sum.631: f32[]) -> f32[] { @@ -1419,61 +1419,61 @@ StackFrames ROOT %reduce_sum.472 = f32[]{:T(128)} add(%reduce_sum.697, %reduce_sum.471), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.515 (param_0.4134: f32[], param_1.5019: f32[256,1,512,512], param_2.4292: f32[], param_3.2945: f32[256,1,512,512], param_4.2197: f32[], param_5.2000: bf16[256,512,512], param_6.1437: pred[], param_7.1118: f32[], param_8.883: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.514 (param_0.4135: f32[], param_1.5015: f32[256,1,512,512], param_2.4290: f32[], param_3.2943: f32[256,1,512,512], param_4.2197: f32[], param_5.2000: bf16[256,512,512], param_6.1437: pred[], param_7.1118: f32[], param_8.883: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.883 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1359.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.883), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %bitcast.1357.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.883), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} %param_7.1118 = f32[]{:T(128)S(6)} parameter(7) - %mul.4676.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4980.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1437 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2147.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1437), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2103.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1437), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2000 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1572.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2000), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1361.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1572.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %broadcast_in_dim.1459.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2000), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1359.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1459.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2197 = f32[]{:T(128)} parameter(4) - %div.2533.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2197), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2532.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1361.clone.1, %div.2533.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2146.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2147.clone.1, %bitcast.1361.clone.1, %div.2532.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4834.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4259.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4834.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} - %mul.4678.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2146.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2945 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1360.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2945), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %constant.4833.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4258.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4833.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4677.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4258.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3408.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4678.clone.1, %mul.4677.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4292 = f32[]{:T(128)S(6)} parameter(2) - %div.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4292), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2146.clone.1, %select_n.2146.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4832.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4261.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4832.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} - %mul.4680.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4261.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5019 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1362.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5019), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %constant.4831.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4260.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4831.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4679.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4260.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3409.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4680.clone.1, %mul.4679.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4134 = f32[]{:T(128)S(6)} parameter(0) - %div.2530.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4134), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2529.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3409.clone.1, %div.2530.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.151.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2529.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4835.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4257.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4835.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3407.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.151.clone.1, %broadcast.4257.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1287.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2531.clone.1, %add.3407.clone.1), metadata={op_name="multiply.296"} - %div.2528.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3408.clone.1, %multiply.1287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4675.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1359.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3406.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2528.clone.1, %mul.4675.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4674.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4676.clone.1, %add.3406.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3405.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1359.clone.1, %mul.4674.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.577 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3405.clone.1, %add.3405.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5069 = f32[]{:T(128)} constant(0) - %reduce.675 = f32[]{:T(128)} reduce(%square.577, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.849.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3409.clone.1) - %bitcast.822.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3408.clone.1) - %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_193.218, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.670 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.675, %add.3405.clone.1, %bitcast.849.clone.1, %bitcast.822.clone.1, %reduce.684.clone.1) + %div.2523.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2197), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2522.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1359.clone.1, %div.2523.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2102.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2103.clone.1, %bitcast.1359.clone.1, %div.2522.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4814.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4232.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4814.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} + %mul.4982.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %broadcast.4232.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2943 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1358.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2943), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %constant.4813.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4231.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4813.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.4981.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1358.clone.1, %broadcast.4231.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4982.clone.1, %mul.4981.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4290 = f32[]{:T(128)S(6)} parameter(2) + %div.2521.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4290), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %select_n.2102.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4812.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4234.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4812.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} + %mul.4984.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4234.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5015 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1360.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5015), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} + %constant.4811.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4233.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4811.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.4983.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4233.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4984.clone.1, %mul.4983.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4135 = f32[]{:T(128)S(6)} parameter(0) + %div.2520.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4135), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2519.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3395.clone.1, %div.2520.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.151.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2519.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4815.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4230.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4815.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %add.3393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.151.clone.1, %broadcast.4230.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1287.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2521.clone.1, %add.3393.clone.1), metadata={op_name="multiply.296"} + %div.2518.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3394.clone.1, %multiply.1287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4979.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1357.clone.1, %broadcast.4232.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3392.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2518.clone.1, %mul.4979.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4978.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4980.clone.1, %add.3392.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3391.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1357.clone.1, %mul.4978.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.331 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3391.clone.1, %add.3391.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5049 = f32[]{:T(128)} constant(0) + %reduce.675 = f32[]{:T(128)} reduce(%square.331, %constant.5049), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.847.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3395.clone.1) + %bitcast.820.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3394.clone.1) + %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5049), dimensions={0,1,2,3}, to_apply=%region_193.218, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.666 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.675, %add.3391.clone.1, %bitcast.847.clone.1, %bitcast.820.clone.1, %reduce.684.clone.1) } %region_226.251 (reduce_sum.928: f32[], reduce_sum.625: f32[]) -> f32[] { @@ -1488,61 +1488,61 @@ StackFrames ROOT %reduce_sum.470 = f32[]{:T(128)} add(%reduce_sum.690, %reduce_sum.465), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.516 (param_0.4135: f32[], param_1.5020: f32[256,1,512,512], param_2.4293: f32[], param_3.2946: f32[256,1,512,512], param_4.2198: f32[], param_5.2001: bf16[256,512,512], param_6.1438: pred[], param_7.1119: f32[], param_8.884: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.515 (param_0.4136: f32[], param_1.5016: f32[256,1,512,512], param_2.4291: f32[], param_3.2944: f32[256,1,512,512], param_4.2198: f32[], param_5.2001: bf16[256,512,512], param_6.1438: pred[], param_7.1119: f32[], param_8.884: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.884 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1363.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.884), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %bitcast.1361.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.884), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} %param_7.1119 = f32[]{:T(128)S(6)} parameter(7) - %mul.4683.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4987.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1438 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2149.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1438), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2105.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1438), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2001 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1573.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2001), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1365.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1573.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %broadcast_in_dim.1460.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2001), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1363.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1460.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2198 = f32[]{:T(128)} parameter(4) - %div.2539.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2198), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2538.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1365.clone.1, %div.2539.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2148.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2149.clone.1, %bitcast.1365.clone.1, %div.2538.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4839.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4264.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4839.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} - %mul.4685.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2148.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2946 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1364.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2946), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %constant.4838.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4263.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4838.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4684.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4263.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3413.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4685.clone.1, %mul.4684.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4293 = f32[]{:T(128)S(6)} parameter(2) - %div.2537.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4293), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2148.clone.1, %select_n.2148.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4837.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4266.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4837.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} - %mul.4687.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4266.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5020 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1366.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5020), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %constant.4836.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4265.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4836.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4686.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4265.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3414.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4687.clone.1, %mul.4686.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4135 = f32[]{:T(128)S(6)} parameter(0) - %div.2536.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4135), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2535.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3414.clone.1, %div.2536.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.152.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2535.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4840.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4262.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3412.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.152.clone.1, %broadcast.4262.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1288.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2537.clone.1, %add.3412.clone.1), metadata={op_name="multiply.295"} - %div.2534.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3413.clone.1, %multiply.1288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4682.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1363.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3411.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2534.clone.1, %mul.4682.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4681.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4683.clone.1, %add.3411.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3410.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1363.clone.1, %mul.4681.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.578 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3410.clone.1, %add.3410.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5070 = f32[]{:T(128)} constant(0) - %reduce.676 = f32[]{:T(128)} reduce(%square.578, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.840.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3414.clone.1) - %bitcast.813.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3413.clone.1) - %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5070), dimensions={0,1,2,3}, to_apply=%region_192.217, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.669 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.676, %add.3410.clone.1, %bitcast.840.clone.1, %bitcast.813.clone.1, %reduce.685.clone.1) + %div.2529.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2198), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2528.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1363.clone.1, %div.2529.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2104.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2105.clone.1, %bitcast.1363.clone.1, %div.2528.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4819.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4237.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4819.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} + %mul.4989.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4237.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2944 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1362.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2944), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %constant.4818.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4236.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4818.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.4988.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4236.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3399.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4989.clone.1, %mul.4988.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4291 = f32[]{:T(128)S(6)} parameter(2) + %div.2527.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4291), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %select_n.2104.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4817.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4239.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4817.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} + %mul.4991.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4239.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5016 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1364.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5016), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} + %constant.4816.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4238.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4816.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.4990.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4238.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3400.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4991.clone.1, %mul.4990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4136 = f32[]{:T(128)S(6)} parameter(0) + %div.2526.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4136), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2525.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3400.clone.1, %div.2526.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.152.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2525.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4820.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4235.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4820.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %add.3398.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.152.clone.1, %broadcast.4235.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1288.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2527.clone.1, %add.3398.clone.1), metadata={op_name="multiply.295"} + %div.2524.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3399.clone.1, %multiply.1288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4986.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1361.clone.1, %broadcast.4237.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3397.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2524.clone.1, %mul.4986.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4985.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4987.clone.1, %add.3397.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3396.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1361.clone.1, %mul.4985.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.332 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3396.clone.1, %add.3396.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5050 = f32[]{:T(128)} constant(0) + %reduce.676 = f32[]{:T(128)} reduce(%square.332, %constant.5050), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.838.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3400.clone.1) + %bitcast.811.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3399.clone.1) + %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5050), dimensions={0,1,2,3}, to_apply=%region_192.217, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.665 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.676, %add.3396.clone.1, %bitcast.838.clone.1, %bitcast.811.clone.1, %reduce.685.clone.1) } %region_225.250 (reduce_sum.921: f32[], reduce_sum.619: f32[]) -> f32[] { @@ -1557,61 +1557,61 @@ StackFrames ROOT %reduce_sum.464 = f32[]{:T(128)} add(%reduce_sum.683, %reduce_sum.463), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.517 (param_0.4136: f32[], param_1.5021: f32[256,1,512,512], param_2.4294: f32[], param_3.2947: f32[256,1,512,512], param_4.2199: f32[], param_5.2002: bf16[256,512,512], param_6.1439: pred[], param_7.1120: f32[], param_8.885: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { +%fused_computation.516 (param_0.4137: f32[], param_1.5017: f32[256,1,512,512], param_2.4292: f32[], param_3.2945: f32[256,1,512,512], param_4.2199: f32[], param_5.2002: bf16[256,512,512], param_6.1439: pred[], param_7.1120: f32[], param_8.885: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { %param_8.885 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1367.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.885), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %bitcast.1365.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.885), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} %param_7.1120 = f32[]{:T(128)S(6)} parameter(7) - %mul.4690.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4994.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1439 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2151.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1439), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2107.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1439), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2002 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1574.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2002), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1369.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1574.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %broadcast_in_dim.1461.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2002), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1367.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1461.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} %param_4.2199 = f32[]{:T(128)} parameter(4) - %div.2545.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2199), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2544.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1369.clone.1, %div.2545.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2150.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2151.clone.1, %bitcast.1369.clone.1, %div.2544.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4844.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4269.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4844.clone.1), dimensions={}, metadata={op_name="broadcast.2345"} - %mul.4692.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2150.clone.1, %broadcast.4269.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2947 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1368.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2947), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4843.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4268.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4843.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4691.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4268.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3418.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4692.clone.1, %mul.4691.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4294 = f32[]{:T(128)S(6)} parameter(2) - %div.2543.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4294), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2150.clone.1, %select_n.2150.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4842.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4271.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4842.clone.1), dimensions={}, metadata={op_name="broadcast.2348"} - %mul.4694.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4271.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5021 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1370.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5021), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4841.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4270.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4693.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1370.clone.1, %broadcast.4270.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3419.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4694.clone.1, %mul.4693.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4136 = f32[]{:T(128)S(6)} parameter(0) - %div.2542.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4136), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2541.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3419.clone.1, %div.2542.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.153.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2541.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4845.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4267.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4845.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3417.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.153.clone.1, %broadcast.4267.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1289.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2543.clone.1, %add.3417.clone.1), metadata={op_name="multiply.294"} - %div.2540.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3418.clone.1, %multiply.1289.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4689.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1367.clone.1, %broadcast.4269.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3416.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2540.clone.1, %mul.4689.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4688.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4690.clone.1, %add.3416.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3415.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1367.clone.1, %mul.4688.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.579 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3415.clone.1, %add.3415.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5071 = f32[]{:T(128)} constant(0) - %reduce.677 = f32[]{:T(128)} reduce(%square.579, %constant.5071), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.831.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3419.clone.1) - %bitcast.804.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3418.clone.1) - %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5071), dimensions={0,1,2,3}, to_apply=%region_191.216, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.668 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.677, %add.3415.clone.1, %bitcast.831.clone.1, %bitcast.804.clone.1, %reduce.686.clone.1) + %div.2535.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2199), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2534.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1367.clone.1, %div.2535.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2106.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2107.clone.1, %bitcast.1367.clone.1, %div.2534.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4824.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4242.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4824.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} + %mul.4996.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %broadcast.4242.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2945 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.1366.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2945), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %constant.4823.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4241.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4823.clone.1), dimensions={}, metadata={op_name="broadcast.329"} + %mul.4995.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4241.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3404.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4996.clone.1, %mul.4995.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4292 = f32[]{:T(128)S(6)} parameter(2) + %div.2533.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4292), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %select_n.2106.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4822.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4244.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4822.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} + %mul.4998.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4244.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5017 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.1368.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5017), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} + %constant.4821.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4243.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4821.clone.1), dimensions={}, metadata={op_name="broadcast.312"} + %mul.4997.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4243.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3405.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4998.clone.1, %mul.4997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4137 = f32[]{:T(128)S(6)} parameter(0) + %div.2532.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4137), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3405.clone.1, %div.2532.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.153.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2531.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4825.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4240.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4825.clone.1), dimensions={}, metadata={op_name="broadcast.305"} + %add.3403.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.153.clone.1, %broadcast.4240.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1289.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2533.clone.1, %add.3403.clone.1), metadata={op_name="multiply.294"} + %div.2530.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3404.clone.1, %multiply.1289.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4993.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1365.clone.1, %broadcast.4242.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3402.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2530.clone.1, %mul.4993.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4992.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4994.clone.1, %add.3402.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3401.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1365.clone.1, %mul.4992.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.333 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3401.clone.1, %add.3401.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5051 = f32[]{:T(128)} constant(0) + %reduce.677 = f32[]{:T(128)} reduce(%square.333, %constant.5051), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.829.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3405.clone.1) + %bitcast.802.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3404.clone.1) + %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5051), dimensions={0,1,2,3}, to_apply=%region_191.216, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.664 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.677, %add.3401.clone.1, %bitcast.829.clone.1, %bitcast.802.clone.1, %reduce.686.clone.1) } %region_155.180 (reduce_sum.438: f32[], reduce_sum.259: f32[]) -> f32[] { @@ -1620,62 +1620,62 @@ StackFrames ROOT %reduce_sum.260 = f32[]{:T(128)} add(%reduce_sum.438, %reduce_sum.259), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.529.clone.clone.clone (param_0.4079: bf16[4,128,129280], param_1.4953: s32[4,128], param_2.4225: f32[4,128], param_3.2913: f32[4,128], param_4.2170: bf16[4,128], param_5.1978: f32[4,128]) -> bf16[4,128,129280] { +%fused_computation.528.clone.clone.clone (param_0.4080: bf16[4,128,129280], param_1.4949: s32[4,128], param_2.4223: f32[4,128], param_3.2911: f32[4,128], param_4.2170: bf16[4,128], param_5.1978: f32[4,128]) -> bf16[4,128,129280] { %param_5.1978 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.4903 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.2913 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.4902 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2913), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.4079 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.3163 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4079), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %mul.5207 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.2911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.5206 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.4080 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.3153 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4080), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} %param_4.2170 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.804 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_4.2170), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.803 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.3163, %sub.804), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.534 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.803), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.4901 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4902, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.4225 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.2698 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4225), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.2697 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.4901, %div.2698), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.4953 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.371 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.4953), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.370 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.369 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.371, %eq.370), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.3162 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%eq.369), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.802 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%div.2697, %convert_element_type.3162), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.4900 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4903, %sub.802), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.3161 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.4900), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.939.clone.clone (param_0.4080: f32[4,128], param_1.4954: bf16[4,128,512], param_2.4227: bf16[512]) -> bf16[4,128,512] { - %param_2.4227 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) - %dot_general.831 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4227), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.4954 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3165 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4954), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.4080 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.4905 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4080), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.4904 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3165, %mul.4905), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.3164 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4904), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.830 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.831, %convert_element_type.3164), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.518 (param_0.4169: bf16[4,128,129280], param_1.5049: s32[4,128], param_2.4319: f32[4,128], param_3.2969: f32[4,128], param_4.2219: bf16[4,128], param_5.2020: f32[4,128], param_6.1457: f32[4,128], param_7.1138: bf16[4,128,512], param_8.902: bf16[512]) -> (f32[], bf16[512,129280,1]) { + %sub.791 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_4.2170), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.790 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.3153, %sub.791), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.534 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.790), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %mul.5205 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.5206, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.4223 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.2688 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4223), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.2687 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.5205, %div.2688), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.4949 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.363 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.4949), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.362 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.361 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.363, %eq.362), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %convert_element_type.3152 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%eq.361), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.789 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%div.2687, %convert_element_type.3152), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.5204 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.5207, %sub.789), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.3151 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.5204), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.938.clone.clone (param_0.4081: f32[4,128], param_1.4950: bf16[4,128,512], param_2.4225: bf16[512]) -> bf16[4,128,512] { + %param_2.4225 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.831 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4225), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.4950 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.3155 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.4081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.5209 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4081), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.5208 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3155, %mul.5209), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.3154 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.5208), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.830 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.831, %convert_element_type.3154), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.517 (param_0.4170: bf16[4,128,129280], param_1.5045: s32[4,128], param_2.4317: f32[4,128], param_3.2967: f32[4,128], param_4.2219: bf16[4,128], param_5.2020: f32[4,128], param_6.1457: f32[4,128], param_7.1138: bf16[4,128,512], param_8.902: bf16[512]) -> (f32[], bf16[512,129280,1]) { %param_6.1457 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) %param_7.1138 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) %param_8.902 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(8) - %fusion.577.clone.1 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_6.1457, %param_7.1138, %param_8.902), kind=kLoop, calls=%fused_computation.939.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.4169 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.5049 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4319 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.2969 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %fusion.574.clone.1 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_6.1457, %param_7.1138, %param_8.902), kind=kLoop, calls=%fused_computation.938.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.4170 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.5045 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.4317 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.2967 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) %param_4.2219 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) %param_5.2020 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} fusion(%param_0.4169, %param_1.5049, %param_2.4319, %param_3.2969, %param_4.2219, /*index=5*/%param_5.2020), kind=kLoop, calls=%fused_computation.529.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.141.clone.1 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.577.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.776 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%convolution.141.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.2665 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.776), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %square.581 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2665, %convert_element_type.2665), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5104 = f32[]{:T(128)} constant(0) - %reduce.678 = f32[]{:T(128)} reduce(%square.581, %constant.5104), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.757 = (f32[]{:T(128)}, bf16[512,129280,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.678, %convolution.141.clone.1) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} fusion(%param_0.4170, %param_1.5045, %param_2.4317, %param_3.2967, %param_4.2219, /*index=5*/%param_5.2020), kind=kLoop, calls=%fused_computation.528.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convolution.141.clone.1 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.574.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %bitcast.774 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%convolution.141.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.2655 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.774), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %mul.3859 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2655, %convert_element_type.2655), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5084 = f32[]{:T(128)} constant(0) + %reduce.678 = f32[]{:T(128)} reduce(%mul.3859, %constant.5084), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.753 = (f32[]{:T(128)}, bf16[512,129280,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.678, %convolution.141.clone.1) } %region_174.199 (reduce_sum.564: f32[], reduce_sum.387: f32[]) -> f32[] { @@ -1684,12 +1684,12 @@ StackFrames ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.564, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.519 (param_0.4153: bf16[129280,512]) -> f32[] { - %param_0.4153 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2667 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4153), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %square.583 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2667, %convert_element_type.2667), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5088 = f32[]{:T(128)} constant(0) - ROOT %reduce.679 = f32[]{:T(128)} reduce(%square.583, %constant.5088), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.518 (param_0.4154: bf16[129280,512]) -> f32[] { + %param_0.4154 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2657 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4154), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %mul.3861 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2657, %convert_element_type.2657), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5068 = f32[]{:T(128)} constant(0) + ROOT %reduce.679 = f32[]{:T(128)} reduce(%mul.3861, %constant.5068), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_240.265 (reduce_sum.1026: f32[], reduce_sum.689: f32[]) -> f32[] { @@ -1704,55 +1704,55 @@ StackFrames ROOT %reduce_sum.534 = f32[]{:T(128)} add(%reduce_sum.788, %reduce_sum.533), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.520 (param_0.4121: f32[129280,512], param_1.5006: f32[], param_2.4279: f32[], param_3.2932: f32[], param_4.2184: f32[129280,512], param_5.1987: f32[], param_6.1424: bf16[129280,512], param_7.1105: pred[], param_8.870: f32[129280,512]) -> (f32[], f32[129280,512], f32[129280,512], f32[129280,512], f32[]) { - %param_0.4121 = f32[129280,512]{1,0:T(8,128)} parameter(0) - %param_3.2932 = f32[]{:T(128)S(6)} parameter(3) - %mul.4564.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2932), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.519 (param_0.4122: f32[129280,512], param_1.5002: f32[], param_2.4277: f32[], param_3.2930: f32[], param_4.2184: f32[129280,512], param_5.1987: f32[], param_6.1424: bf16[129280,512], param_7.1105: pred[], param_8.870: f32[129280,512]) -> (f32[], f32[129280,512], f32[129280,512], f32[129280,512], f32[]) { + %param_0.4122 = f32[129280,512]{1,0:T(8,128)} parameter(0) + %param_3.2930 = f32[]{:T(128)S(6)} parameter(3) + %mul.4868.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1105 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2105.clone.1 = pred[129280,512]{1,0:T(8,128)(4,1)} broadcast(%param_7.1105), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2061.clone.1 = pred[129280,512]{1,0:T(8,128)(4,1)} broadcast(%param_7.1105), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1424 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.3106.clone.1 = f32[129280,512]{1,0:T(8,128)} convert(%param_6.1424), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convert_element_type.3096.clone.1 = f32[129280,512]{1,0:T(8,128)} convert(%param_6.1424), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %param_5.1987 = f32[]{:T(128)} parameter(5) - %div.2439.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_5.1987), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2438.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%convert_element_type.3106.clone.1, %div.2439.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2104.clone.1 = f32[129280,512]{1,0:T(8,128)} select(%select_n.2105.clone.1, %convert_element_type.3106.clone.1, %div.2438.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4754.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4209.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4754.clone.1), dimensions={}, metadata={op_name="broadcast.318"} - %mul.4570.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4209.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2429.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_5.1987), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2428.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%convert_element_type.3096.clone.1, %div.2429.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2060.clone.1 = f32[129280,512]{1,0:T(8,128)} select(%select_n.2061.clone.1, %convert_element_type.3096.clone.1, %div.2428.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4734.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4182.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4734.clone.1), dimensions={}, metadata={op_name="broadcast.318"} + %mul.4874.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %broadcast.4182.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.870 = f32[129280,512]{1,0:T(8,128)} parameter(8) - %constant.4758.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4571.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4758.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4569.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3338.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4570.clone.1, %mul.4569.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4279 = f32[]{:T(128)S(6)} parameter(2) - %div.2435.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4279), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.380.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2104.clone.1, %select_n.2104.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4757.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4568.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4757.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4566.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4568.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4738.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4875.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4738.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4873.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3324.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4874.clone.1, %mul.4873.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4277 = f32[]{:T(128)S(6)} parameter(2) + %div.2425.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.380.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %select_n.2060.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4737.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4872.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4737.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4870.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2184 = f32[129280,512]{1,0:T(8,128)} parameter(4) - %constant.4756.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4567.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4756.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4565.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4567.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3337.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4566.clone.1, %mul.4565.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5006 = f32[]{:T(128)S(6)} parameter(1) - %div.2434.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2433.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3337.clone.1, %div.2434.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.138.clone.1 = f32[129280,512]{1,0:T(8,128)} sqrt(%div.2433.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4755.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3336.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4755.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3335.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%sqrt.138.clone.1, %add.3336.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1274.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%div.2435.clone.1, %add.3335.clone.1), metadata={op_name="multiply.309"} - %div.2432.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3338.clone.1, %multiply.1274.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4563.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4121, %broadcast.4209.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3334.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2432.clone.1, %mul.4563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4562.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4564.clone.1, %add.3334.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3333.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4121, %mul.4562.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.584 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3333.clone.1, %add.3333.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5056 = f32[]{:T(128)} constant(0) - %reduce.680 = f32[]{:T(128)} reduce(%square.584, %constant.5056), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.687.clone.1 = f32[]{:T(128)} reduce(%integer_pow.380.clone.1, %constant.5056), dimensions={0,1}, to_apply=%region_206.231, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.671 = (f32[]{:T(128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.680, %add.3333.clone.1, %add.3337.clone.1, %add.3338.clone.1, %reduce.687.clone.1) + %constant.4736.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4871.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4736.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4869.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3323.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4870.clone.1, %mul.4869.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5002 = f32[]{:T(128)S(6)} parameter(1) + %div.2424.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5002), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2423.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3323.clone.1, %div.2424.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.138.clone.1 = f32[129280,512]{1,0:T(8,128)} sqrt(%div.2423.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4735.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3322.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4735.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3321.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%sqrt.138.clone.1, %add.3322.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1274.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%div.2425.clone.1, %add.3321.clone.1), metadata={op_name="multiply.309"} + %div.2422.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3324.clone.1, %multiply.1274.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4867.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4122, %broadcast.4182.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3320.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2422.clone.1, %mul.4867.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4866.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4868.clone.1, %add.3320.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3319.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4122, %mul.4866.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.334 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3319.clone.1, %add.3319.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5036 = f32[]{:T(128)} constant(0) + %reduce.680 = f32[]{:T(128)} reduce(%square.334, %constant.5036), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.687.clone.1 = f32[]{:T(128)} reduce(%integer_pow.380.clone.1, %constant.5036), dimensions={0,1}, to_apply=%region_206.231, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.667 = (f32[]{:T(128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.680, %add.3319.clone.1, %add.3323.clone.1, %add.3324.clone.1, %reduce.687.clone.1) } %region_222.247 (reduce_sum.900: f32[], reduce_sum.605: f32[]) -> f32[] { @@ -1767,56 +1767,56 @@ StackFrames ROOT %reduce_sum.455 = f32[]{:T(128)} add(%reduce_sum.662, %reduce_sum.451), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.521 (param_0.4139: f32[512,129280], param_1.5024: f32[], param_2.4297: f32[], param_3.2950: f32[], param_4.2202: f32[512,129280], param_5.2005: f32[], param_6.1442: bf16[512,129280,1], param_7.1123: pred[], param_8.888: f32[512,129280]) -> (f32[], f32[512,129280], f32[512,129280], f32[512,129280], f32[]) { - %param_0.4139 = f32[512,129280]{1,0:T(8,128)} parameter(0) - %param_3.2950 = f32[]{:T(128)S(6)} parameter(3) - %mul.4717.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2950), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.520 (param_0.4140: f32[512,129280], param_1.5020: f32[], param_2.4295: f32[], param_3.2948: f32[], param_4.2202: f32[512,129280], param_5.2005: f32[], param_6.1442: bf16[512,129280,1], param_7.1123: pred[], param_8.888: f32[512,129280]) -> (f32[], f32[512,129280], f32[512,129280], f32[512,129280], f32[]) { + %param_0.4140 = f32[512,129280]{1,0:T(8,128)} parameter(0) + %param_3.2948 = f32[]{:T(128)S(6)} parameter(3) + %mul.5021.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2948), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1123 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2161.clone.1 = pred[512,129280]{1,0:T(8,128)(4,1)} broadcast(%param_7.1123), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2117.clone.1 = pred[512,129280]{1,0:T(8,128)(4,1)} broadcast(%param_7.1123), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1442 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.1372.clone.1 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%param_6.1442), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.3108.clone.1 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.1372.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %bitcast.1370.clone.1 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%param_6.1442), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.3098.clone.1 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.1370.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} %param_5.2005 = f32[]{:T(128)} parameter(5) - %div.2567.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_5.2005), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2566.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%convert_element_type.3108.clone.1, %div.2567.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2160.clone.1 = f32[512,129280]{1,0:T(8,128)} select(%select_n.2161.clone.1, %convert_element_type.3108.clone.1, %div.2566.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4858.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4277.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="broadcast.333"} - %mul.4723.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2160.clone.1, %broadcast.4277.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2557.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_5.2005), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2556.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%convert_element_type.3098.clone.1, %div.2557.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2116.clone.1 = f32[512,129280]{1,0:T(8,128)} select(%select_n.2117.clone.1, %convert_element_type.3098.clone.1, %div.2556.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4838.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4250.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4838.clone.1), dimensions={}, metadata={op_name="broadcast.333"} + %mul.5027.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.888 = f32[512,129280]{1,0:T(8,128)} parameter(8) - %constant.4862.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4724.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4862.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4722.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.4724.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3437.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4723.clone.1, %mul.4722.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4297 = f32[]{:T(128)S(6)} parameter(2) - %div.2563.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4297), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.398.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2160.clone.1, %select_n.2160.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4861.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4721.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4861.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4719.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.4721.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4842.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.5028.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4842.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5026.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.5028.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3423.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.5027.clone.1, %mul.5026.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4295 = f32[]{:T(128)S(6)} parameter(2) + %div.2553.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4295), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.398.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %select_n.2116.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4841.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.5025.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5023.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.5025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2202 = f32[512,129280]{1,0:T(8,128)} parameter(4) - %constant.4860.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4720.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4718.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.4720.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3436.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4719.clone.1, %mul.4718.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5024 = f32[]{:T(128)S(6)} parameter(1) - %div.2562.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5024), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2561.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3436.clone.1, %div.2562.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.156.clone.1 = f32[512,129280]{1,0:T(8,128)} sqrt(%div.2561.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4859.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3435.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4859.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3434.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%sqrt.156.clone.1, %add.3435.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1292.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%div.2563.clone.1, %add.3434.clone.1), metadata={op_name="multiply.291"} - %div.2560.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3437.clone.1, %multiply.1292.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4716.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4139, %broadcast.4277.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3433.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2560.clone.1, %mul.4716.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4715.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.4717.clone.1, %add.3433.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3432.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4139, %mul.4715.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.585 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3432.clone.1, %add.3432.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5074 = f32[]{:T(128)} constant(0) - %reduce.681 = f32[]{:T(128)} reduce(%square.585, %constant.5074), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.688.clone.1 = f32[]{:T(128)} reduce(%integer_pow.398.clone.1, %constant.5074), dimensions={0,1}, to_apply=%region_188.213, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.672 = (f32[]{:T(128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.681, %add.3432.clone.1, %add.3436.clone.1, %add.3437.clone.1, %reduce.688.clone.1) + %constant.4840.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.5024.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5022.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.5024.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3422.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.5023.clone.1, %mul.5022.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5020 = f32[]{:T(128)S(6)} parameter(1) + %div.2552.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5020), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2551.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3422.clone.1, %div.2552.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.156.clone.1 = f32[512,129280]{1,0:T(8,128)} sqrt(%div.2551.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4839.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3421.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4839.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3420.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%sqrt.156.clone.1, %add.3421.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1292.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%div.2553.clone.1, %add.3420.clone.1), metadata={op_name="multiply.291"} + %div.2550.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3423.clone.1, %multiply.1292.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.5020.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4140, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3419.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2550.clone.1, %mul.5020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.5019.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.5021.clone.1, %add.3419.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3418.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4140, %mul.5019.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.335 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3418.clone.1, %add.3418.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5054 = f32[]{:T(128)} constant(0) + %reduce.681 = f32[]{:T(128)} reduce(%square.335, %constant.5054), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.688.clone.1 = f32[]{:T(128)} reduce(%integer_pow.398.clone.1, %constant.5054), dimensions={0,1}, to_apply=%region_188.213, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.668 = (f32[]{:T(128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.681, %add.3418.clone.1, %add.3422.clone.1, %add.3423.clone.1, %reduce.688.clone.1) } %region_207.232 (reduce_sum.795: f32[], reduce_sum.535: f32[]) -> f32[] { @@ -1825,23 +1825,23 @@ StackFrames ROOT %reduce_sum.540 = f32[]{:T(128)} add(%reduce_sum.795, %reduce_sum.535), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.522 (param_0.4190: bf16[4,128,129280], param_1.5063: f32[4,128], param_2.4329: s32[4,128], param_3.2977: bf16[4,128]) -> f32[4,128] { - %param_2.4329 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.307 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4329), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.302 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.301 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.307, %eq.302), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.4190 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2672 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4190), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.2977 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.665 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2977), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.656 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2672, %sub.665), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.5063 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.663 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5063), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.652 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%sub.656, %sub.663), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.5128 = f32[]{:T(128)} constant(0) - %broadcast.3784 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5128), dimensions={}, metadata={op_name="broadcast.518"} - %mul.3624 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.301, %sub.652, %broadcast.3784), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3624, %constant.5128), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +%fused_computation.521 (param_0.4191: bf16[4,128,129280], param_1.5059: f32[4,128], param_2.4327: s32[4,128], param_3.2975: bf16[4,128]) -> f32[4,128] { + %param_2.4327 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.299 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4327), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.294 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.293 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.299, %eq.294), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %param_0.4191 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2662 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_3.2975 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) + %sub.652 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2975), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.643 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2662, %sub.652), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.5059 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.650 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5059), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.639 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%sub.643, %sub.650), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %constant.5108 = f32[]{:T(128)} constant(0) + %broadcast.3757 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5108), dimensions={}, metadata={op_name="broadcast.514"} + %mul.3862 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.293, %sub.639, %broadcast.3757), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3862, %constant.5108), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_37.47 (reduce_sum.76: f32[], reduce_sum.80: f32[]) -> f32[] { @@ -1850,15 +1850,15 @@ StackFrames ROOT %reduce_sum.83 = f32[]{:T(128)} add(%reduce_sum.76, %reduce_sum.80), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.533 (param_0.4191: bf16[4,128,129280], param_1.5064: bf16[4,128]) -> f32[4,128] { - %param_0.4191 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2678 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.5064 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.666 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5064), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.662 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2678, %sub.666), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.448 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.662), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.5129 = f32[]{:T(128)} constant(0) - ROOT %reduce.683 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.448, %constant.5129), dimensions={2}, to_apply=%region_37.47, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +%fused_computation.532 (param_0.4192: bf16[4,128,129280], param_1.5060: bf16[4,128]) -> f32[4,128] { + %param_0.4192 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.2668 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.5060 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.653 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5060), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.649 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2668, %sub.653), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.448 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.649), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %constant.5109 = f32[]{:T(128)} constant(0) + ROOT %reduce.683 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.448, %constant.5109), dimensions={2}, to_apply=%region_37.47, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_152.177 (reduce_sum.417: f32[], reduce_sum.244: f32[]) -> f32[] { @@ -1867,18 +1867,18 @@ StackFrames ROOT %reduce_sum.251 = f32[]{:T(128)} add(%reduce_sum.417, %reduce_sum.244), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.541 (param_0.4172: f32[3,512,128,256]) -> f32[] { - %param_0.4172 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.752 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4172), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.588 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.752, %bitcast.752), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5107 = f32[]{:T(128)} constant(0) - ROOT %reduce.689 = f32[]{:T(128)} reduce(%square.588, %constant.5107), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +%fused_computation.540 (param_0.4173: f32[3,512,128,256]) -> f32[] { + %param_0.4173 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.750 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4173), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3883 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.750, %bitcast.750), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5087 = f32[]{:T(128)} constant(0) + ROOT %reduce.689 = f32[]{:T(128)} reduce(%mul.3883, %constant.5087), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.542 (param_0.1602: f32[512,3,128,256]) -> bf16[3,512,128,256] { - %param_0.1602 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) - %copy.1551 = bf16[512,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1602), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\']"} - ROOT %bitcast.753 = bf16[3,512,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1551), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} +%fused_computation.541 (param_0.1600: f32[512,3,128,256]) -> bf16[3,512,128,256] { + %param_0.1600 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) + %copy.1551 = bf16[512,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1600), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\']"} + ROOT %bitcast.751 = bf16[3,512,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1551), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} } %region_219.244 (reduce_sum.879: f32[], reduce_sum.591: f32[]) -> f32[] { @@ -1893,55 +1893,55 @@ StackFrames ROOT %reduce_sum.442 = f32[]{:T(128)} add(%reduce_sum.641, %reduce_sum.437), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.543 (param_0.4142: f32[512,3,128,256], param_1.5027: f32[], param_2.4300: f32[], param_3.2953: f32[], param_4.2205: f32[512,3,128,256], param_5.2008: f32[], param_6.1445: f32[3,512,128,256], param_7.1126: pred[], param_8.891: f32[512,3,128,256]) -> (f32[], f32[512,3,128,256], f32[512,3,128,256], f32[512,3,128,256], f32[]) { - %param_0.4142 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) - %param_3.2953 = f32[]{:T(128)S(6)} parameter(3) - %mul.4747.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2953), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} +%fused_computation.542 (param_0.4143: f32[512,3,128,256], param_1.5023: f32[], param_2.4298: f32[], param_3.2951: f32[], param_4.2205: f32[512,3,128,256], param_5.2008: f32[], param_6.1445: f32[3,512,128,256], param_7.1126: pred[], param_8.891: f32[512,3,128,256]) -> (f32[], f32[512,3,128,256], f32[512,3,128,256], f32[512,3,128,256], f32[]) { + %param_0.4143 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) + %param_3.2951 = f32[]{:T(128)S(6)} parameter(3) + %mul.5051.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1126 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2173.clone.1 = pred[512,3,128,256]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.1126), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %select_n.2129.clone.1 = pred[512,3,128,256]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.1126), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1445 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.1378.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_6.1445), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} + %bitcast.1376.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_6.1445), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} %param_5.2008 = f32[]{:T(128)} parameter(5) - %div.2591.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_5.2008), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2590.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.1378.clone.1, %div.2591.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2172.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2173.clone.1, %bitcast.1378.clone.1, %div.2590.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4876.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4283.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4876.clone.1), dimensions={}, metadata={op_name="broadcast.336"} - %mul.4753.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2172.clone.1, %broadcast.4283.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %div.2581.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_5.2008), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2580.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.1376.clone.1, %div.2581.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2128.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2129.clone.1, %bitcast.1376.clone.1, %div.2580.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.4856.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4256.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4856.clone.1), dimensions={}, metadata={op_name="broadcast.336"} + %mul.5057.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.891 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(8) - %constant.4880.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4754.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4880.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4752.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.4754.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3455.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4753.clone.1, %mul.4752.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4300 = f32[]{:T(128)S(6)} parameter(2) - %div.2587.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4300), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.401.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2172.clone.1, %select_n.2172.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4879.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4751.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4879.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4749.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.4751.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.4860.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.5058.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5056.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.5058.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3441.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.5057.clone.1, %mul.5056.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4298 = f32[]{:T(128)S(6)} parameter(2) + %div.2577.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4298), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.401.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %select_n.2128.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.4859.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.5055.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4859.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5053.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.5055.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2205 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(4) - %constant.4878.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4750.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4878.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4748.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.4750.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3454.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4749.clone.1, %mul.4748.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5027 = f32[]{:T(128)S(6)} parameter(1) - %div.2586.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5027), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2585.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3454.clone.1, %div.2586.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.159.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} sqrt(%div.2585.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4877.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3453.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4877.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3452.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3453.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1295.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%div.2587.clone.1, %add.3452.clone.1), metadata={op_name="multiply.288"} - %div.2584.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3455.clone.1, %multiply.1295.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4746.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4142, %broadcast.4283.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3451.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2584.clone.1, %mul.4746.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4745.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.4747.clone.1, %add.3451.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3450.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4142, %mul.4745.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.589 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3450.clone.1, %add.3450.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5077 = f32[]{:T(128)} constant(0) - %reduce.690 = f32[]{:T(128)} reduce(%square.589, %constant.5077), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.691.clone.1 = f32[]{:T(128)} reduce(%integer_pow.401.clone.1, %constant.5077), dimensions={0,1,2,3}, to_apply=%region_185.210, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.667 = (f32[]{:T(128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.690, %add.3450.clone.1, %add.3454.clone.1, %add.3455.clone.1, %reduce.691.clone.1) + %constant.4858.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.5054.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.5052.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.5054.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3440.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.5053.clone.1, %mul.5052.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5023 = f32[]{:T(128)S(6)} parameter(1) + %div.2576.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5023), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2575.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3440.clone.1, %div.2576.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.159.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} sqrt(%div.2575.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.4857.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3439.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4857.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3438.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3439.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1295.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%div.2577.clone.1, %add.3438.clone.1), metadata={op_name="multiply.288"} + %div.2574.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3441.clone.1, %multiply.1295.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.5050.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4143, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3437.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2574.clone.1, %mul.5050.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.5049.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.5051.clone.1, %add.3437.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3436.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4143, %mul.5049.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.336 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3436.clone.1, %add.3436.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.5057 = f32[]{:T(128)} constant(0) + %reduce.690 = f32[]{:T(128)} reduce(%square.336, %constant.5057), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.691.clone.1 = f32[]{:T(128)} reduce(%integer_pow.401.clone.1, %constant.5057), dimensions={0,1,2,3}, to_apply=%region_185.210, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.663 = (f32[]{:T(128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.690, %add.3436.clone.1, %add.3440.clone.1, %add.3441.clone.1, %reduce.691.clone.1) } %region_172.197 (reduce_sum.557: f32[], reduce_sum.381: f32[]) -> f32[] { @@ -1950,39 +1950,39 @@ StackFrames ROOT %reduce_sum.386 = f32[]{:T(128)} add(%reduce_sum.557, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.783.clone.clone (param_0.4106: f32[4,128], param_1.4998: bf16[4,128,1536], param_2.4261: bf16[1536]) -> bf16[4,128,1536,1] { - %param_2.4261 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.851 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4261), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.4998 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3187 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4998), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} - %param_0.4106 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.4951 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4106), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %mul.4950 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3187, %mul.4951), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %convert_element_type.3186 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.4950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} - %dot_general.850 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.851, %convert_element_type.3186), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.1466 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dot_general.850), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.782.clone.clone (param_0.4107: f32[4,128], param_1.4994: bf16[4,128,1536], param_2.4259: bf16[1536]) -> bf16[4,128,1536,1] { + %param_2.4259 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.851 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4259), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.4994 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.3177 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4994), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} + %param_0.4107 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.5255 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4107), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} + %mul.5254 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3177, %mul.5255), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} + %convert_element_type.3176 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.5254), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} + %dot_general.850 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.851, %convert_element_type.3176), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.1464 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dot_general.850), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} } %bitcast_fusion.12 (bitcast_input.12: bf16[4,128,128,192]) -> bf16[4,128,128,192] { %bitcast_input.12 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.1488 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} bitcast(%bitcast_input.12) -} - -%fused_computation.552 (param_0.4154: bf16[4,128,128,192], param_1.5038: f32[4,128], param_2.4311: bf16[4,128,1536], param_3.2964: bf16[1536]) -> (f32[], bf16[1536,128,192,1]) { - %param_1.5038 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4311 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.2964 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.460.clone.1 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.5038, %param_2.4311, %param_3.2964), kind=kLoop, calls=%fused_computation.783.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.4154 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - %fusion.751 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.4154), kind=kLoop, calls=%bitcast_fusion.12 - %convolution.146.clone.1 = bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)} convolution(%fusion.460.clone.1, %fusion.751), window={size=192x4 pad=191_191x0_0 rhs_reversal=1x0}, dim_labels=1fb0_1io0->bf01, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} - %bitcast.861 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%convolution.146.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} - %broadcast_in_dim.1388 = f32[1536,128,192]{1,0,2:T(8,128)} convert(%bitcast.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.763 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.592 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.763, %bitcast.763), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5089 = f32[]{:T(128)} constant(0) - %reduce.692 = f32[]{:T(128)} reduce(%square.592, %constant.5089), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.766 = (f32[]{:T(128)}, bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)}) tuple(%reduce.692, %convolution.146.clone.1) + ROOT %bitcast.1486 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} bitcast(%bitcast_input.12) +} + +%fused_computation.551 (param_0.4155: bf16[4,128,128,192], param_1.5034: f32[4,128], param_2.4309: bf16[4,128,1536], param_3.2962: bf16[1536]) -> (f32[], bf16[1536,128,192,1]) { + %param_1.5034 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.4309 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.2962 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.457.clone.1 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.5034, %param_2.4309, %param_3.2962), kind=kLoop, calls=%fused_computation.782.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.4155 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) + %fusion.748 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.4155), kind=kLoop, calls=%bitcast_fusion.12 + %convolution.144.clone.1 = bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)} convolution(%fusion.457.clone.1, %fusion.748), window={size=192x4 pad=191_191x0_0 rhs_reversal=1x0}, dim_labels=1fb0_1io0->bf01, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} + %bitcast.859 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%convolution.144.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} + %broadcast_in_dim.1275 = f32[1536,128,192]{1,0,2:T(8,128)} convert(%bitcast.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.761 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1275), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} + %mul.3892 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.761, %bitcast.761), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.5069 = f32[]{:T(128)} constant(0) + %reduce.692 = f32[]{:T(128)} reduce(%mul.3892, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.762 = (f32[]{:T(128)}, bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)}) tuple(%reduce.692, %convolution.144.clone.1) } %region_239.264 (reduce_sum.1019: f32[], reduce_sum.687: f32[]) -> f32[] { @@ -1997,4 +1997,4 @@ StackFrames ROOT %reduce_sum.528 = f32[]{:T(128)} add(%reduce_sum.781, %reduce_sum.527), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.557 (param_0.4122: f32[], param_1.5007: f32[], param_2.4280: f32[], param_3.2933: f32[1536,1,128,192], param_4.2185: f32[1536,1,128,192], param_5.1988: f32[], param_6.1425: bf16[1536,128,192,1], param_7.1106: pred[], param_8.871: f32[1536,1,128,192]) -> (f32[], f32[1536,1,128,192], f32[1536,1,128,192], f32[1536,1,128,192], f32[]) { +%fused_computation.556 (param_0.4123: f32[], param_1.5003: f32[], param_2.4278: f32[], param_3.2931: f32[1536,1,128,192], param_4.2185: f32[1536,1,128,192], param_5.1988: f32[], param_6.1425: bf16[1536,128,192,1], param_7.1106: pred[], param_8.871: f32[1536,1,128,192]) -> (f32[], f32[1536,1,128,192], f32[1536,1,128,192], f32[1536,1,128,192], f32[]) { From ac3259965545c05f5a1a6ee24c200e344c1019fb Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Wed, 10 Jun 2026 02:20:15 +0000 Subject: [PATCH 28/52] Fix QwixQuantization.einsum signature and revert workflow changes - Updated QwixQuantization.einsum to accept **kwargs to prevent TypeError when called from MoE layer with dtype. - Reverted .github/workflows/run_tests_against_package.yml to main to isolate whether the NCCL failures were triggered by the new workflow changes (e.g., NCCL_SOCKET_IFNAME=lo). TAG=agy CONV=753fe72c-2db8-4329-9321-b25762bed269 --- .../workflows/run_tests_against_package.yml | 30 ++++++------------- src/maxtext/layers/quantizations.py | 2 +- 2 files changed, 10 insertions(+), 22 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index abc76932db..9df299ab21 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -162,28 +162,16 @@ jobs: # Dynamically discover the 'nvidia' folder and prepend all its sub-library # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent # JAX from partially loading incompatible system-level CUDA libraries. - NVIDIA_DIR=$(find -L .venv -path "*/site-packages/nvidia" -type d 2>/dev/null | head -n 1) - echo "=== GPU Diagnostics ===" - echo "Found NVIDIA_DIR: ${NVIDIA_DIR}" - if [ -n "${NVIDIA_DIR}" ]; then - for dir in "${NVIDIA_DIR}"/*; do - if [ -d "$dir/lib" ]; then - ABS_LIB_PATH=$(realpath "$dir/lib") - export LD_LIBRARY_PATH=${ABS_LIB_PATH}:${LD_LIBRARY_PATH} - echo "Prepended to LD_LIBRARY_PATH: ${ABS_LIB_PATH}" - fi - done - else - echo "WARNING: nvidia directory not found under .venv!" + if [ -d ".venv/lib" ]; then + NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1) + if [ -n "${NVIDIA_DIR}" ]; then + for dir in "${NVIDIA_DIR}"/*; do + if [ -d "$dir/lib" ]; then + export LD_LIBRARY_PATH=$(pwd)/$dir/lib:${LD_LIBRARY_PATH} + fi + done + fi fi - echo "Final LD_LIBRARY_PATH: ${LD_LIBRARY_PATH}" - echo "=======================" - - # Configure NCCL for stable single-node communication in Docker containers - export NCCL_SOCKET_IFNAME=lo - export NCCL_NET_GDR_LEVEL=0 - export NCCL_DEBUG=INFO - echo "Set NCCL_SOCKET_IFNAME=lo, NCCL_NET_GDR_LEVEL=0, and NCCL_DEBUG=INFO for GPU execution." fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index f7f64f7385..3f3af1bf06 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -81,7 +81,7 @@ def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Returns Qwix dot_general.""" return functools.partial(QwixDotGeneral, config=self._get_fp8_full_qwix_config()) - def einsum(self, mesh_axes: Tuple[str, ...] = ()): + def einsum(self, mesh_axes: Tuple[str, ...] = (), **kwargs): """Returns Qwix einsum.""" return QwixEinsum(config=self._get_fp8_full_qwix_config()) From beba1d1d72875b743326f27efaea0b5fae582788 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Wed, 10 Jun 2026 15:49:01 +0000 Subject: [PATCH 29/52] Restore GPU diagnostics and NCCL config in workflow --- .../workflows/run_tests_against_package.yml | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 9df299ab21..abc76932db 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -162,16 +162,28 @@ jobs: # Dynamically discover the 'nvidia' folder and prepend all its sub-library # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent # JAX from partially loading incompatible system-level CUDA libraries. - if [ -d ".venv/lib" ]; then - NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1) - if [ -n "${NVIDIA_DIR}" ]; then - for dir in "${NVIDIA_DIR}"/*; do - if [ -d "$dir/lib" ]; then - export LD_LIBRARY_PATH=$(pwd)/$dir/lib:${LD_LIBRARY_PATH} - fi - done - fi + NVIDIA_DIR=$(find -L .venv -path "*/site-packages/nvidia" -type d 2>/dev/null | head -n 1) + echo "=== GPU Diagnostics ===" + echo "Found NVIDIA_DIR: ${NVIDIA_DIR}" + if [ -n "${NVIDIA_DIR}" ]; then + for dir in "${NVIDIA_DIR}"/*; do + if [ -d "$dir/lib" ]; then + ABS_LIB_PATH=$(realpath "$dir/lib") + export LD_LIBRARY_PATH=${ABS_LIB_PATH}:${LD_LIBRARY_PATH} + echo "Prepended to LD_LIBRARY_PATH: ${ABS_LIB_PATH}" + fi + done + else + echo "WARNING: nvidia directory not found under .venv!" fi + echo "Final LD_LIBRARY_PATH: ${LD_LIBRARY_PATH}" + echo "=======================" + + # Configure NCCL for stable single-node communication in Docker containers + export NCCL_SOCKET_IFNAME=lo + export NCCL_NET_GDR_LEVEL=0 + export NCCL_DEBUG=INFO + echo "Set NCCL_SOCKET_IFNAME=lo, NCCL_NET_GDR_LEVEL=0, and NCCL_DEBUG=INFO for GPU execution." fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist From 009209e865d8dac6a57eacfee55416cd44b25ed9 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Wed, 10 Jun 2026 17:09:54 +0000 Subject: [PATCH 30/52] Update HLO references to match CI environment (Flax 0.12.7, Optax 0.2.8) --- tests/utils/reference_hlo_deepseek3.txt | 266 +-- tests/utils/reference_hlo_llama3_8b.txt | 2086 +++++++++++----------- tests/utils/reference_hlo_qwen3_1.7b.txt | 1906 ++++++++++---------- 3 files changed, 2129 insertions(+), 2129 deletions(-) diff --git a/tests/utils/reference_hlo_deepseek3.txt b/tests/utils/reference_hlo_deepseek3.txt index cc3471559c..5413169ee5 100644 --- a/tests/utils/reference_hlo_deepseek3.txt +++ b/tests/utils/reference_hlo_deepseek3.txt @@ -1291,9 +1291,9 @@ StackFrames %fused_computation.466 (param_0.4171: f32[3,1536,128,192]) -> f32[] { %param_0.4171 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(0) %bitcast.670 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4171), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3786 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.670, %bitcast.670), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %square.564 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.670, %bitcast.670), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5085 = f32[]{:T(128)} constant(0) - ROOT %reduce.669 = f32[]{:T(128)} reduce(%mul.3786, %constant.5085), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.669 = f32[]{:T(128)} reduce(%square.564, %constant.5085), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %fused_computation.467 (param_0.1419: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { @@ -1317,7 +1317,7 @@ StackFrames %fused_computation.468 (param_0.4141: f32[1536,3,128,192], param_1.5021: f32[], param_2.4296: f32[], param_3.2949: f32[], param_4.2203: f32[1536,3,128,192], param_5.2006: f32[], param_6.1443: f32[3,1536,128,192], param_7.1124: pred[], param_8.889: f32[1536,3,128,192]) -> (f32[], f32[1536,3,128,192], f32[1536,3,128,192], f32[1536,3,128,192], f32[]) { %param_0.4141 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) %param_3.2949 = f32[]{:T(128)S(6)} parameter(3) - %mul.5031.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2949), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4715.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2949), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1124 = pred[]{:T(512)S(6)} parameter(7) %select_n.2121.clone.1 = pred[1536,3,128,192]{2,3,1,0:T(8,128)(4,1)} broadcast(%param_7.1124), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1443 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(6) @@ -1328,23 +1328,23 @@ StackFrames %select_n.2120.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2121.clone.1, %bitcast.1372.clone.1, %div.2564.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.4844.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.4252.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4844.clone.1), dimensions={}, metadata={op_name="broadcast.334"} - %mul.5037.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4721.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.889 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(8) %constant.4848.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.5038.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4848.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5036.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.5038.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3429.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.5037.clone.1, %mul.5036.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4722.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4848.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4720.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.4722.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3429.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4721.clone.1, %mul.4720.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_2.4296 = f32[]{:T(128)S(6)} parameter(2) %div.2561.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4296), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.399.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %select_n.2120.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.4847.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.5035.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4847.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5033.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.5035.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4719.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4847.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4717.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.4719.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2203 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(4) %constant.4846.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.5034.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4846.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5032.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.5034.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3428.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.5033.clone.1, %mul.5032.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4718.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4846.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4716.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.4718.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3428.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4717.clone.1, %mul.4716.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_1.5021 = f32[]{:T(128)S(6)} parameter(1) %div.2560.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5021), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2559.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3428.clone.1, %div.2560.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} @@ -1354,13 +1354,13 @@ StackFrames %add.3426.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.157.clone.1, %add.3427.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1293.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%div.2561.clone.1, %add.3426.clone.1), metadata={op_name="multiply.290"} %div.2558.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3429.clone.1, %multiply.1293.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.5030.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4141, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3425.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2558.clone.1, %mul.5030.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.5029.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.5031.clone.1, %add.3425.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3424.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4141, %mul.5029.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.330 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3424.clone.1, %add.3424.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.4714.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4141, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3425.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2558.clone.1, %mul.4714.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4713.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.4715.clone.1, %add.3425.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3424.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4141, %mul.4713.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.565 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3424.clone.1, %add.3424.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5055 = f32[]{:T(128)} constant(0) - %reduce.670 = f32[]{:T(128)} reduce(%square.330, %constant.5055), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.670 = f32[]{:T(128)} reduce(%square.565, %constant.5055), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.671.clone.1 = f32[]{:T(128)} reduce(%integer_pow.399.clone.1, %constant.5055), dimensions={0,1,2,3}, to_apply=%region_187.212, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.656 = (f32[]{:T(128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.670, %add.3424.clone.1, %add.3428.clone.1, %add.3429.clone.1, %reduce.671.clone.1) } @@ -1381,14 +1381,14 @@ StackFrames %param_0.4167 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) %broadcast_in_dim.1245 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4167), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} %bitcast.693 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3815 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.693, %bitcast.693), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %square.570 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.693, %bitcast.693), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5081 = f32[]{:T(128)} constant(0) - %reduce.672 = f32[]{:T(128)} reduce(%mul.3815, %constant.5081), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.672 = f32[]{:T(128)} reduce(%square.570, %constant.5081), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %param_1.5043 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(1) %broadcast_in_dim.1253.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_1.5043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} %bitcast.701.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1253.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3821.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.701.clone.1, %bitcast.701.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %reduce.674.clone.1 = f32[]{:T(128)} reduce(%mul.3821.clone.1, %constant.5081), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %square.576.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.701.clone.1, %bitcast.701.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.674.clone.1 = f32[]{:T(128)} reduce(%square.576.clone.1, %constant.5081), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.763 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.672, %reduce.674.clone.1) } @@ -1402,9 +1402,9 @@ StackFrames %param_0.4166 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) %broadcast_in_dim.1249 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4166), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} %bitcast.697 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3818 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.697, %bitcast.697), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %square.573 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.697, %bitcast.697), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5080 = f32[]{:T(128)} constant(0) - ROOT %reduce.673 = f32[]{:T(128)} reduce(%mul.3818, %constant.5080), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.673 = f32[]{:T(128)} reduce(%square.573, %constant.5080), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_227.252 (reduce_sum.935: f32[], reduce_sum.631: f32[]) -> f32[] { @@ -1423,7 +1423,7 @@ StackFrames %param_8.883 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) %bitcast.1357.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.883), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} %param_7.1118 = f32[]{:T(128)S(6)} parameter(7) - %mul.4980.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4664.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1437 = pred[]{:T(512)S(6)} parameter(6) %select_n.2103.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1437), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2000 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) @@ -1435,25 +1435,25 @@ StackFrames %select_n.2102.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2103.clone.1, %bitcast.1359.clone.1, %div.2522.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.4814.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.4232.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4814.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} - %mul.4982.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %broadcast.4232.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4666.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %broadcast.4232.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_3.2943 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) %bitcast.1358.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2943), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} %constant.4813.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.4231.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4813.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4981.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1358.clone.1, %broadcast.4231.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4982.clone.1, %mul.4981.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4665.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1358.clone.1, %broadcast.4231.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4666.clone.1, %mul.4665.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_2.4290 = f32[]{:T(128)S(6)} parameter(2) %div.2521.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4290), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %select_n.2102.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.4812.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.4234.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4812.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} - %mul.4984.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4234.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4668.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4234.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_1.5015 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) %bitcast.1360.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5015), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} %constant.4811.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.4233.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4811.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4983.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4233.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4984.clone.1, %mul.4983.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4667.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4233.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4668.clone.1, %mul.4667.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_0.4135 = f32[]{:T(128)S(6)} parameter(0) %div.2520.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4135), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2519.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3395.clone.1, %div.2520.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} @@ -1463,13 +1463,13 @@ StackFrames %add.3393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.151.clone.1, %broadcast.4230.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1287.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2521.clone.1, %add.3393.clone.1), metadata={op_name="multiply.296"} %div.2518.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3394.clone.1, %multiply.1287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4979.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1357.clone.1, %broadcast.4232.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3392.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2518.clone.1, %mul.4979.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4978.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4980.clone.1, %add.3392.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3391.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1357.clone.1, %mul.4978.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.331 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3391.clone.1, %add.3391.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.4663.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1357.clone.1, %broadcast.4232.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3392.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2518.clone.1, %mul.4663.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4662.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4664.clone.1, %add.3392.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3391.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1357.clone.1, %mul.4662.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.577 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3391.clone.1, %add.3391.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5049 = f32[]{:T(128)} constant(0) - %reduce.675 = f32[]{:T(128)} reduce(%square.331, %constant.5049), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.675 = f32[]{:T(128)} reduce(%square.577, %constant.5049), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %bitcast.847.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3395.clone.1) %bitcast.820.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3394.clone.1) %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5049), dimensions={0,1,2,3}, to_apply=%region_193.218, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1492,7 +1492,7 @@ StackFrames %param_8.884 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) %bitcast.1361.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.884), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} %param_7.1119 = f32[]{:T(128)S(6)} parameter(7) - %mul.4987.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4671.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1438 = pred[]{:T(512)S(6)} parameter(6) %select_n.2105.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1438), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2001 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) @@ -1504,25 +1504,25 @@ StackFrames %select_n.2104.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2105.clone.1, %bitcast.1363.clone.1, %div.2528.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.4819.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.4237.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4819.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} - %mul.4989.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4237.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4673.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4237.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_3.2944 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) %bitcast.1362.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2944), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} %constant.4818.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.4236.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4818.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4988.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4236.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3399.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4989.clone.1, %mul.4988.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4672.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4236.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3399.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4673.clone.1, %mul.4672.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_2.4291 = f32[]{:T(128)S(6)} parameter(2) %div.2527.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4291), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %select_n.2104.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.4817.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.4239.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4817.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} - %mul.4991.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4239.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4675.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4239.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_1.5016 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) %bitcast.1364.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5016), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} %constant.4816.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.4238.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4816.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4990.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4238.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3400.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4991.clone.1, %mul.4990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4674.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4238.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3400.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4675.clone.1, %mul.4674.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_0.4136 = f32[]{:T(128)S(6)} parameter(0) %div.2526.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4136), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2525.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3400.clone.1, %div.2526.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} @@ -1532,13 +1532,13 @@ StackFrames %add.3398.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.152.clone.1, %broadcast.4235.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1288.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2527.clone.1, %add.3398.clone.1), metadata={op_name="multiply.295"} %div.2524.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3399.clone.1, %multiply.1288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4986.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1361.clone.1, %broadcast.4237.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3397.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2524.clone.1, %mul.4986.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4985.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4987.clone.1, %add.3397.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3396.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1361.clone.1, %mul.4985.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.332 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3396.clone.1, %add.3396.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.4670.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1361.clone.1, %broadcast.4237.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3397.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2524.clone.1, %mul.4670.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4669.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4671.clone.1, %add.3397.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3396.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1361.clone.1, %mul.4669.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.578 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3396.clone.1, %add.3396.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5050 = f32[]{:T(128)} constant(0) - %reduce.676 = f32[]{:T(128)} reduce(%square.332, %constant.5050), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.676 = f32[]{:T(128)} reduce(%square.578, %constant.5050), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %bitcast.838.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3400.clone.1) %bitcast.811.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3399.clone.1) %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5050), dimensions={0,1,2,3}, to_apply=%region_192.217, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1561,7 +1561,7 @@ StackFrames %param_8.885 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) %bitcast.1365.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.885), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} %param_7.1120 = f32[]{:T(128)S(6)} parameter(7) - %mul.4994.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4678.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_6.1439 = pred[]{:T(512)S(6)} parameter(6) %select_n.2107.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1439), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_5.2002 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) @@ -1573,25 +1573,25 @@ StackFrames %select_n.2106.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2107.clone.1, %bitcast.1367.clone.1, %div.2534.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.4824.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.4242.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4824.clone.1), dimensions={}, metadata={op_name="broadcast.2344"} - %mul.4996.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %broadcast.4242.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4680.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %broadcast.4242.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_3.2945 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) %bitcast.1366.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2945), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} %constant.4823.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.4241.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4823.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4995.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4241.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3404.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4996.clone.1, %mul.4995.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4679.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4241.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3404.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4680.clone.1, %mul.4679.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_2.4292 = f32[]{:T(128)S(6)} parameter(2) %div.2533.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4292), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %select_n.2106.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.4822.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.4244.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4822.clone.1), dimensions={}, metadata={op_name="broadcast.2347"} - %mul.4998.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4244.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4682.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4244.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_1.5017 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) %bitcast.1368.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5017), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} %constant.4821.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.4243.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4821.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4997.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4243.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3405.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4998.clone.1, %mul.4997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4681.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4243.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3405.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4682.clone.1, %mul.4681.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_0.4137 = f32[]{:T(128)S(6)} parameter(0) %div.2532.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4137), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3405.clone.1, %div.2532.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} @@ -1601,13 +1601,13 @@ StackFrames %add.3403.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.153.clone.1, %broadcast.4240.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1289.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2533.clone.1, %add.3403.clone.1), metadata={op_name="multiply.294"} %div.2530.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3404.clone.1, %multiply.1289.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4993.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1365.clone.1, %broadcast.4242.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3402.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2530.clone.1, %mul.4993.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4992.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4994.clone.1, %add.3402.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3401.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1365.clone.1, %mul.4992.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.333 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3401.clone.1, %add.3401.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.4677.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1365.clone.1, %broadcast.4242.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3402.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2530.clone.1, %mul.4677.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4676.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4678.clone.1, %add.3402.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3401.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1365.clone.1, %mul.4676.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.579 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3401.clone.1, %add.3401.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5051 = f32[]{:T(128)} constant(0) - %reduce.677 = f32[]{:T(128)} reduce(%square.333, %constant.5051), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.677 = f32[]{:T(128)} reduce(%square.579, %constant.5051), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %bitcast.829.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3405.clone.1) %bitcast.802.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3404.clone.1) %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5051), dimensions={0,1,2,3}, to_apply=%region_191.216, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1622,27 +1622,27 @@ StackFrames %fused_computation.528.clone.clone.clone (param_0.4080: bf16[4,128,129280], param_1.4949: s32[4,128], param_2.4223: f32[4,128], param_3.2911: f32[4,128], param_4.2170: bf16[4,128], param_5.1978: f32[4,128]) -> bf16[4,128,129280] { %param_5.1978 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.5207 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.4891 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.2911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.5206 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.4890 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_0.4080 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) %convert_element_type.3153 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4080), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} %param_4.2170 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) %sub.791 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_4.2170), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.790 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.3153, %sub.791), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.534 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.790), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.5205 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.5206, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.4889 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4890, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_2.4223 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) %div.2688 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4223), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.2687 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.5205, %div.2688), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.2687 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.4889, %div.2688), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} %param_1.4949 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) %eq.363 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.4949), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.362 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.361 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.363, %eq.362), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %convert_element_type.3152 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%eq.361), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} %sub.789 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%div.2687, %convert_element_type.3152), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.5204 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.5207, %sub.789), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.3151 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.5204), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %mul.4888 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4891, %sub.789), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.3151 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.4888), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} } %fused_computation.938.clone.clone (param_0.4081: f32[4,128], param_1.4950: bf16[4,128,512], param_2.4225: bf16[512]) -> bf16[4,128,512] { @@ -1651,9 +1651,9 @@ StackFrames %param_1.4950 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) %convert_element_type.3155 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %param_0.4081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.5209 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4081), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.5208 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3155, %mul.5209), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.3154 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.5208), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.4893 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4081), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.4892 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3155, %mul.4893), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.3154 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4892), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} ROOT %dot_general.830 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.831, %convert_element_type.3154), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } @@ -1672,9 +1672,9 @@ StackFrames %convolution.141.clone.1 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.574.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} %bitcast.774 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%convolution.141.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} %convert_element_type.2655 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.774), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %mul.3859 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2655, %convert_element_type.2655), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %square.581 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2655, %convert_element_type.2655), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5084 = f32[]{:T(128)} constant(0) - %reduce.678 = f32[]{:T(128)} reduce(%mul.3859, %constant.5084), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.678 = f32[]{:T(128)} reduce(%square.581, %constant.5084), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.753 = (f32[]{:T(128)}, bf16[512,129280,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.678, %convolution.141.clone.1) } @@ -1687,9 +1687,9 @@ StackFrames %fused_computation.518 (param_0.4154: bf16[129280,512]) -> f32[] { %param_0.4154 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) %convert_element_type.2657 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4154), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %mul.3861 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2657, %convert_element_type.2657), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %square.583 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2657, %convert_element_type.2657), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5068 = f32[]{:T(128)} constant(0) - ROOT %reduce.679 = f32[]{:T(128)} reduce(%mul.3861, %constant.5068), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.679 = f32[]{:T(128)} reduce(%square.583, %constant.5068), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_240.265 (reduce_sum.1026: f32[], reduce_sum.689: f32[]) -> f32[] { @@ -1707,7 +1707,7 @@ StackFrames %fused_computation.519 (param_0.4122: f32[129280,512], param_1.5002: f32[], param_2.4277: f32[], param_3.2930: f32[], param_4.2184: f32[129280,512], param_5.1987: f32[], param_6.1424: bf16[129280,512], param_7.1105: pred[], param_8.870: f32[129280,512]) -> (f32[], f32[129280,512], f32[129280,512], f32[129280,512], f32[]) { %param_0.4122 = f32[129280,512]{1,0:T(8,128)} parameter(0) %param_3.2930 = f32[]{:T(128)S(6)} parameter(3) - %mul.4868.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4552.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1105 = pred[]{:T(512)S(6)} parameter(7) %select_n.2061.clone.1 = pred[129280,512]{1,0:T(8,128)(4,1)} broadcast(%param_7.1105), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1424 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(6) @@ -1718,23 +1718,23 @@ StackFrames %select_n.2060.clone.1 = f32[129280,512]{1,0:T(8,128)} select(%select_n.2061.clone.1, %convert_element_type.3096.clone.1, %div.2428.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.4734.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.4182.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4734.clone.1), dimensions={}, metadata={op_name="broadcast.318"} - %mul.4874.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %broadcast.4182.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4558.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %broadcast.4182.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.870 = f32[129280,512]{1,0:T(8,128)} parameter(8) %constant.4738.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4875.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4738.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4873.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3324.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4874.clone.1, %mul.4873.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4559.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4738.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4557.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3324.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4558.clone.1, %mul.4557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_2.4277 = f32[]{:T(128)S(6)} parameter(2) %div.2425.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.380.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %select_n.2060.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.4737.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4872.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4737.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4870.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4556.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4737.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4554.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2184 = f32[129280,512]{1,0:T(8,128)} parameter(4) %constant.4736.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4871.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4736.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4869.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3323.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4870.clone.1, %mul.4869.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4555.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4736.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4553.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4555.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3323.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4554.clone.1, %mul.4553.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_1.5002 = f32[]{:T(128)S(6)} parameter(1) %div.2424.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5002), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2423.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3323.clone.1, %div.2424.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} @@ -1744,13 +1744,13 @@ StackFrames %add.3321.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%sqrt.138.clone.1, %add.3322.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1274.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%div.2425.clone.1, %add.3321.clone.1), metadata={op_name="multiply.309"} %div.2422.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3324.clone.1, %multiply.1274.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4867.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4122, %broadcast.4182.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3320.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2422.clone.1, %mul.4867.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4866.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4868.clone.1, %add.3320.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3319.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4122, %mul.4866.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.334 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3319.clone.1, %add.3319.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.4551.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4122, %broadcast.4182.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3320.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2422.clone.1, %mul.4551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4550.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4552.clone.1, %add.3320.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3319.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4122, %mul.4550.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.584 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3319.clone.1, %add.3319.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5036 = f32[]{:T(128)} constant(0) - %reduce.680 = f32[]{:T(128)} reduce(%square.334, %constant.5036), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.680 = f32[]{:T(128)} reduce(%square.584, %constant.5036), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.687.clone.1 = f32[]{:T(128)} reduce(%integer_pow.380.clone.1, %constant.5036), dimensions={0,1}, to_apply=%region_206.231, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.667 = (f32[]{:T(128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.680, %add.3319.clone.1, %add.3323.clone.1, %add.3324.clone.1, %reduce.687.clone.1) } @@ -1770,7 +1770,7 @@ StackFrames %fused_computation.520 (param_0.4140: f32[512,129280], param_1.5020: f32[], param_2.4295: f32[], param_3.2948: f32[], param_4.2202: f32[512,129280], param_5.2005: f32[], param_6.1442: bf16[512,129280,1], param_7.1123: pred[], param_8.888: f32[512,129280]) -> (f32[], f32[512,129280], f32[512,129280], f32[512,129280], f32[]) { %param_0.4140 = f32[512,129280]{1,0:T(8,128)} parameter(0) %param_3.2948 = f32[]{:T(128)S(6)} parameter(3) - %mul.5021.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2948), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4705.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2948), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1123 = pred[]{:T(512)S(6)} parameter(7) %select_n.2117.clone.1 = pred[512,129280]{1,0:T(8,128)(4,1)} broadcast(%param_7.1123), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1442 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} parameter(6) @@ -1782,23 +1782,23 @@ StackFrames %select_n.2116.clone.1 = f32[512,129280]{1,0:T(8,128)} select(%select_n.2117.clone.1, %convert_element_type.3098.clone.1, %div.2556.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.4838.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.4250.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4838.clone.1), dimensions={}, metadata={op_name="broadcast.333"} - %mul.5027.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4711.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.888 = f32[512,129280]{1,0:T(8,128)} parameter(8) %constant.4842.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.5028.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4842.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5026.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.5028.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3423.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.5027.clone.1, %mul.5026.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4712.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4842.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4710.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.4712.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3423.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4711.clone.1, %mul.4710.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_2.4295 = f32[]{:T(128)S(6)} parameter(2) %div.2553.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4295), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.398.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %select_n.2116.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.4841.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.5025.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5023.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.5025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4709.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4707.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.4709.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2202 = f32[512,129280]{1,0:T(8,128)} parameter(4) %constant.4840.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.5024.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5022.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.5024.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3422.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.5023.clone.1, %mul.5022.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4708.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4706.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.4708.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3422.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4707.clone.1, %mul.4706.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_1.5020 = f32[]{:T(128)S(6)} parameter(1) %div.2552.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5020), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2551.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3422.clone.1, %div.2552.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} @@ -1808,13 +1808,13 @@ StackFrames %add.3420.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%sqrt.156.clone.1, %add.3421.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1292.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%div.2553.clone.1, %add.3420.clone.1), metadata={op_name="multiply.291"} %div.2550.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3423.clone.1, %multiply.1292.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.5020.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4140, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3419.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2550.clone.1, %mul.5020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.5019.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.5021.clone.1, %add.3419.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3418.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4140, %mul.5019.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.335 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3418.clone.1, %add.3418.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.4704.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4140, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3419.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2550.clone.1, %mul.4704.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4703.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.4705.clone.1, %add.3419.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3418.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4140, %mul.4703.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.585 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3418.clone.1, %add.3418.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5054 = f32[]{:T(128)} constant(0) - %reduce.681 = f32[]{:T(128)} reduce(%square.335, %constant.5054), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.681 = f32[]{:T(128)} reduce(%square.585, %constant.5054), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.688.clone.1 = f32[]{:T(128)} reduce(%integer_pow.398.clone.1, %constant.5054), dimensions={0,1}, to_apply=%region_188.213, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.668 = (f32[]{:T(128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.681, %add.3418.clone.1, %add.3422.clone.1, %add.3423.clone.1, %reduce.688.clone.1) } @@ -1840,8 +1840,8 @@ StackFrames %sub.639 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%sub.643, %sub.650), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %constant.5108 = f32[]{:T(128)} constant(0) %broadcast.3757 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5108), dimensions={}, metadata={op_name="broadcast.514"} - %mul.3862 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.293, %sub.639, %broadcast.3757), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3862, %constant.5108), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.3612 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.293, %sub.639, %broadcast.3757), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3612, %constant.5108), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_37.47 (reduce_sum.76: f32[], reduce_sum.80: f32[]) -> f32[] { @@ -1870,9 +1870,9 @@ StackFrames %fused_computation.540 (param_0.4173: f32[3,512,128,256]) -> f32[] { %param_0.4173 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(0) %bitcast.750 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4173), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3883 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.750, %bitcast.750), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %square.588 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.750, %bitcast.750), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5087 = f32[]{:T(128)} constant(0) - ROOT %reduce.689 = f32[]{:T(128)} reduce(%mul.3883, %constant.5087), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %reduce.689 = f32[]{:T(128)} reduce(%square.588, %constant.5087), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %fused_computation.541 (param_0.1600: f32[512,3,128,256]) -> bf16[3,512,128,256] { @@ -1896,7 +1896,7 @@ StackFrames %fused_computation.542 (param_0.4143: f32[512,3,128,256], param_1.5023: f32[], param_2.4298: f32[], param_3.2951: f32[], param_4.2205: f32[512,3,128,256], param_5.2008: f32[], param_6.1445: f32[3,512,128,256], param_7.1126: pred[], param_8.891: f32[512,3,128,256]) -> (f32[], f32[512,3,128,256], f32[512,3,128,256], f32[512,3,128,256], f32[]) { %param_0.4143 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) %param_3.2951 = f32[]{:T(128)S(6)} parameter(3) - %mul.5051.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4735.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_7.1126 = pred[]{:T(512)S(6)} parameter(7) %select_n.2129.clone.1 = pred[512,3,128,256]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.1126), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %param_6.1445 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(6) @@ -1907,23 +1907,23 @@ StackFrames %select_n.2128.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2129.clone.1, %bitcast.1376.clone.1, %div.2580.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.4856.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.4256.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4856.clone.1), dimensions={}, metadata={op_name="broadcast.336"} - %mul.5057.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4741.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_8.891 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(8) %constant.4860.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.5058.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5056.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.5058.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3441.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.5057.clone.1, %mul.5056.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4742.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4740.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.4742.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3441.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4741.clone.1, %mul.4740.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_2.4298 = f32[]{:T(128)S(6)} parameter(2) %div.2577.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4298), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.401.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %select_n.2128.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.4859.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.5055.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4859.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5053.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.5055.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4739.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4859.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4737.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.4739.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %param_4.2205 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(4) %constant.4858.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.5054.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.5052.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.5054.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3440.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.5053.clone.1, %mul.5052.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4738.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4736.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.4738.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3440.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4737.clone.1, %mul.4736.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %param_1.5023 = f32[]{:T(128)S(6)} parameter(1) %div.2576.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5023), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.2575.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3440.clone.1, %div.2576.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} @@ -1933,13 +1933,13 @@ StackFrames %add.3438.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3439.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.1295.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%div.2577.clone.1, %add.3438.clone.1), metadata={op_name="multiply.288"} %div.2574.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3441.clone.1, %multiply.1295.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.5050.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4143, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3437.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2574.clone.1, %mul.5050.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.5049.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.5051.clone.1, %add.3437.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3436.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4143, %mul.5049.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.336 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3436.clone.1, %add.3436.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %mul.4734.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4143, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3437.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2574.clone.1, %mul.4734.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4733.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.4735.clone.1, %add.3437.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3436.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4143, %mul.4733.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.589 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3436.clone.1, %add.3436.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5057 = f32[]{:T(128)} constant(0) - %reduce.690 = f32[]{:T(128)} reduce(%square.336, %constant.5057), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.690 = f32[]{:T(128)} reduce(%square.589, %constant.5057), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.691.clone.1 = f32[]{:T(128)} reduce(%integer_pow.401.clone.1, %constant.5057), dimensions={0,1,2,3}, to_apply=%region_185.210, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.663 = (f32[]{:T(128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.690, %add.3436.clone.1, %add.3440.clone.1, %add.3441.clone.1, %reduce.691.clone.1) } @@ -1956,9 +1956,9 @@ StackFrames %param_1.4994 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) %convert_element_type.3177 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4994), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} %param_0.4107 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.5255 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4107), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %mul.5254 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3177, %mul.5255), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %convert_element_type.3176 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.5254), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} + %mul.4939 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4107), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} + %mul.4938 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3177, %mul.4939), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} + %convert_element_type.3176 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.4938), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} %dot_general.850 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.851, %convert_element_type.3176), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} ROOT %bitcast.1464 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dot_general.850), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} } @@ -1979,9 +1979,9 @@ StackFrames %bitcast.859 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%convolution.144.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} %broadcast_in_dim.1275 = f32[1536,128,192]{1,0,2:T(8,128)} convert(%bitcast.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} %bitcast.761 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1275), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %mul.3892 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.761, %bitcast.761), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %square.592 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.761, %bitcast.761), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.5069 = f32[]{:T(128)} constant(0) - %reduce.692 = f32[]{:T(128)} reduce(%mul.3892, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.692 = f32[]{:T(128)} reduce(%square.592, %constant.5069), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.762 = (f32[]{:T(128)}, bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)}) tuple(%reduce.692, %convolution.144.clone.1) } diff --git a/tests/utils/reference_hlo_llama3_8b.txt b/tests/utils/reference_hlo_llama3_8b.txt index 27c6529df2..488affcd35 100644 --- a/tests/utils/reference_hlo_llama3_8b.txt +++ b/tests/utils/reference_hlo_llama3_8b.txt @@ -44,62 +44,62 @@ StackFrames ROOT %reduce_sum.192 = f32[]{:T(128)} add(%reduce_sum.190, %reduce_sum.191), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.280.clone.clone.clone (param_0.1099: bf16[4,128,128256], param_1.1265: s32[4,128], param_2.1086: f32[4,128], param_3.785: f32[4,128], param_4.487: bf16[4,128], param_5.412: f32[4,128]) -> bf16[4,128,128256] { - %param_5.412 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1613 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.785 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1612 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.785), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1099 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1044 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1099), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.487 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.487), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1044, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.281.clone.clone.clone (param_0.1085: bf16[4,128,128256], param_1.1251: s32[4,128], param_2.1077: f32[4,128], param_3.781: f32[4,128], param_4.482: bf16[4,128], param_5.404: f32[4,128]) -> bf16[4,128,128256] { + %param_5.404 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.1607 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.404), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.781 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.1606 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.781), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1085 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1032 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1085), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.482 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.482), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1032, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.62 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1611 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1612, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1611, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1265 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1265), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.1605 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1606, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1077 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1077), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1605, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1251 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1251), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.48 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.47 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1043 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1043), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1610 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1613, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1042 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1610), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.316.clone.clone (param_0.1100: f32[4,128], param_1.1266: bf16[4,128,4096], param_2.1088: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1088 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.387 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1088), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1266 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1046 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1266), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1100 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1615 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1100), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1614 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1046, %mul.1615), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1045 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1614), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.387, %convert_element_type.1045), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.219 (param_0.1119: bf16[4,128,128256], param_1.1281: s32[4,128], param_2.1112: f32[4,128], param_3.801: f32[4,128], param_4.502: bf16[4,128], param_5.427: f32[4,128], param_6.299: f32[4,128], param_7.198: bf16[4,128,4096], param_8.116: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { - %param_6.299 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.198 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) - %param_8.116 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) - %fusion.239.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.299, %param_7.198, %param_8.116), kind=kLoop, calls=%fused_computation.316.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1119 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1281 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1112 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.801 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %param_4.502 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %param_5.427 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1119, %param_1.1281, %param_2.1112, %param_3.801, %param_4.502, /*index=5*/%param_5.427), kind=kLoop, calls=%fused_computation.280.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.88.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.239.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.306 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.88.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.923 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.306), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %square.157 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.923, %convert_element_type.923), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %convert_element_type.1031 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1031), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1604 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1607, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1030 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1604), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.317.clone.clone (param_0.1086: f32[4,128], param_1.1252: bf16[4,128,4096], param_2.1079: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1079 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1079), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1034 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1252), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1609 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1608 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1034, %mul.1609), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1033 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1608), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.220 (param_0.1105: bf16[4,128,128256], param_1.1267: s32[4,128], param_2.1103: f32[4,128], param_3.797: f32[4,128], param_4.497: bf16[4,128], param_5.419: f32[4,128], param_6.287: f32[4,128], param_7.186: bf16[4,128,4096], param_8.112: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { + %param_6.287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.186 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) + %param_8.112 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) + %fusion.229.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.287, %param_7.186, %param_8.112), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1105 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1267 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1103 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.797 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %param_4.497 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %param_5.419 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1105, %param_1.1267, %param_2.1103, %param_3.797, %param_4.497, /*index=5*/%param_5.419), kind=kLoop, calls=%fused_computation.281.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convolution.82.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.229.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %bitcast.300 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.82.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.911 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.300), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %square.157 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.911, %convert_element_type.911), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1006 = f32[]{:T(128)} constant(0) %reduce.118 = f32[]{:T(128)} reduce(%square.157, %constant.1006), dimensions={0,1}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.88.clone.1) + ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.82.clone.1) } %region_34.39 (reduce_sum.196: f32[], reduce_sum.197: f32[]) -> f32[] { @@ -108,10 +108,10 @@ StackFrames ROOT %reduce_sum.198 = f32[]{:T(128)} add(%reduce_sum.196, %reduce_sum.197), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.220 (param_0.1118: bf16[128256,4096]) -> f32[] { - %param_0.1118 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.925 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1118), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %square.159 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.925, %convert_element_type.925), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.221 (param_0.1104: bf16[128256,4096]) -> f32[] { + %param_0.1104 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.913 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1104), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %square.159 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.913, %convert_element_type.913), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1005 = f32[]{:T(128)} constant(0) ROOT %reduce.119 = f32[]{:T(128)} reduce(%square.159, %constant.1005), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } @@ -128,39 +128,39 @@ StackFrames ROOT %reduce_sum.261 = f32[]{:T(128)} add(%reduce_sum.259, %reduce_sum.260), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.221 (param_0.1106: f32[128256,4096], param_1.1269: f32[], param_2.1100: f32[], param_3.789: f32[], param_4.490: f32[128256,4096], param_5.415: f32[], param_6.287: bf16[128256,4096], param_7.186: pred[], param_8.104: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { - %param_0.1106 = f32[128256,4096]{1,0:T(8,128)} parameter(0) - %param_3.789 = f32[]{:T(128)S(6)} parameter(3) - %mul.1482.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.186 = pred[]{:T(512)S(6)} parameter(7) - %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.186), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.287 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1017.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.287), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %param_5.415 = f32[]{:T(128)} parameter(5) - %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1017.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1017.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.222 (param_0.1092: f32[128256,4096], param_1.1255: f32[], param_2.1091: f32[], param_3.785: f32[], param_4.485: f32[128256,4096], param_5.407: f32[], param_6.275: bf16[128256,4096], param_7.174: pred[], param_8.100: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { + %param_0.1092 = f32[128256,4096]{1,0:T(8,128)} parameter(0) + %param_3.785 = f32[]{:T(128)S(6)} parameter(3) + %mul.1476.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.785), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.174 = pred[]{:T(512)S(6)} parameter(7) + %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.174), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.275 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1005.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.275), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_5.407 = f32[]{:T(128)} parameter(5) + %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.407), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1005.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1005.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.907.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.554.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.907.clone.1), dimensions={}, metadata={op_name="broadcast.61"} - %mul.1488.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.104 = f32[128256,4096]{1,0:T(8,128)} parameter(8) + %mul.1482.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.100 = f32[128256,4096]{1,0:T(8,128)} parameter(8) %constant.911.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1489.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1487.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.104, %mul.1489.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1488.clone.1, %mul.1487.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) - %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1483.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1481.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.100, %mul.1483.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1482.clone.1, %mul.1481.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1091 = f32[]{:T(128)S(6)} parameter(2) + %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1091), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.60.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %select_n.241.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.910.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1486.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1484.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1486.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.490 = f32[128256,4096]{1,0:T(8,128)} parameter(4) + %mul.1480.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1478.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1480.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.485 = f32[128256,4096]{1,0:T(8,128)} parameter(4) %constant.909.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1485.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1483.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.490, %mul.1485.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1484.clone.1, %mul.1483.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1269 = f32[]{:T(128)S(6)} parameter(1) - %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1269), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1479.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1477.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.485, %mul.1479.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1478.clone.1, %mul.1477.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1255 = f32[]{:T(128)S(6)} parameter(1) + %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1255), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.719.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.775.clone.1, %div.720.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.58.clone.1 = f32[128256,4096]{1,0:T(8,128)} sqrt(%div.719.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.908.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -168,10 +168,10 @@ StackFrames %add.773.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%sqrt.58.clone.1, %add.774.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.256.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%div.721.clone.1, %add.773.clone.1), metadata={op_name="multiply.42"} %div.718.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.776.clone.1, %multiply.256.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1481.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1106, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1481.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1480.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1482.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1106, %mul.1480.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1475.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1092, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1475.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1474.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1476.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1092, %mul.1474.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.160 = f32[128256,4096]{1,0:T(8,128)} multiply(%add.771.clone.1, %add.771.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.993 = f32[]{:T(128)} constant(0) %reduce.120 = f32[]{:T(128)} reduce(%square.160, %constant.993), dimensions={0,1}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -191,40 +191,40 @@ StackFrames ROOT %reduce_sum.255 = f32[]{:T(128)} add(%reduce_sum.253, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.222 (param_0.1107: f32[4096,128256], param_1.1270: f32[], param_2.1101: f32[], param_3.790: f32[], param_4.491: f32[4096,128256], param_5.416: f32[], param_6.288: bf16[4096,128256,1], param_7.187: pred[], param_8.105: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { - %param_0.1107 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %param_3.790 = f32[]{:T(128)S(6)} parameter(3) - %mul.1492.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.187 = pred[]{:T(512)S(6)} parameter(7) - %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.187), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.288 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.409.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.288), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.1019.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.409.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %param_5.416 = f32[]{:T(128)} parameter(5) - %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1019.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1019.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.223 (param_0.1093: f32[4096,128256], param_1.1256: f32[], param_2.1092: f32[], param_3.786: f32[], param_4.486: f32[4096,128256], param_5.408: f32[], param_6.276: bf16[4096,128256,1], param_7.175: pred[], param_8.101: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { + %param_0.1093 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %param_3.786 = f32[]{:T(128)S(6)} parameter(3) + %mul.1486.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.786), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.175 = pred[]{:T(512)S(6)} parameter(7) + %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.175), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.276 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) + %bitcast.403.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.1007.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.403.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %param_5.408 = f32[]{:T(128)} parameter(5) + %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.408), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1007.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1007.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.913.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.556.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.913.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %mul.1498.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.105 = f32[4096,128256]{1,0:T(8,128)} parameter(8) + %mul.1492.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.101 = f32[4096,128256]{1,0:T(8,128)} parameter(8) %constant.917.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1499.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1497.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.105, %mul.1499.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1498.clone.1, %mul.1497.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) - %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1493.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1491.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.101, %mul.1493.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1492.clone.1, %mul.1491.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1092 = f32[]{:T(128)S(6)} parameter(2) + %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1092), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.61.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %select_n.245.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.916.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1496.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1494.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1496.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.491 = f32[4096,128256]{1,0:T(8,128)} parameter(4) + %mul.1490.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1488.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1490.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.486 = f32[4096,128256]{1,0:T(8,128)} parameter(4) %constant.915.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1495.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1493.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.491, %mul.1495.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1494.clone.1, %mul.1493.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1270 = f32[]{:T(128)S(6)} parameter(1) - %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1270), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1489.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1487.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.486, %mul.1489.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1488.clone.1, %mul.1487.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1256 = f32[]{:T(128)S(6)} parameter(1) + %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1256), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.727.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.781.clone.1, %div.728.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.59.clone.1 = f32[4096,128256]{1,0:T(8,128)} sqrt(%div.727.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.914.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -232,10 +232,10 @@ StackFrames %add.779.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%sqrt.59.clone.1, %add.780.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.257.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%div.729.clone.1, %add.779.clone.1), metadata={op_name="multiply.41"} %div.726.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.782.clone.1, %multiply.257.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1491.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1107, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1491.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1490.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1492.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1107, %mul.1490.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1485.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1093, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1485.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1484.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1486.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1093, %mul.1484.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.161 = f32[4096,128256]{1,0:T(8,128)} multiply(%add.777.clone.1, %add.777.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.994 = f32[]{:T(128)} constant(0) %reduce.121 = f32[]{:T(128)} reduce(%square.161, %constant.994), dimensions={0,1}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -249,10 +249,10 @@ StackFrames ROOT %reduce_sum.156 = f32[]{:T(128)} add(%reduce_sum.154, %reduce_sum.155), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.239 (param_0.1124: f32[4,14336,4096]) -> f32[] { - %param_0.1124 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) - %bitcast.314 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1124), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.164 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.314, %bitcast.314), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.240 (param_0.1110: f32[4,14336,4096]) -> f32[] { + %param_0.1110 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) + %bitcast.308 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.164 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.308, %bitcast.308), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1011 = f32[]{:T(128)} constant(0) ROOT %reduce.124 = f32[]{:T(128)} reduce(%square.164, %constant.1011), dimensions={0,1,2}, to_apply=%region_25.30, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } @@ -269,35 +269,35 @@ StackFrames ROOT %reduce_sum.147 = f32[]{:T(128)} add(%reduce_sum.142, %reduce_sum.143), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.241 (param_0.1125: f32[4,4096,14336], param_1.1284: f32[4,4096,14336]) -> (f32[], f32[]) { - %param_0.1125 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) - %bitcast.318 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1125), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.167 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.318, %bitcast.318), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.242 (param_0.1111: f32[4,4096,14336], param_1.1270: f32[4,4096,14336]) -> (f32[], f32[]) { + %param_0.1111 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) + %bitcast.312 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.167 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.312, %bitcast.312), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1012 = f32[]{:T(128)} constant(0) %reduce.125 = f32[]{:T(128)} reduce(%square.167, %constant.1012), dimensions={0,1,2}, to_apply=%region_24.29, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) - %bitcast.322.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.170.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.322.clone.1, %bitcast.322.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %param_1.1270 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) + %bitcast.316.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1270), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.170.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.316.clone.1, %bitcast.316.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %reduce.126.clone.1 = f32[]{:T(128)} reduce(%square.170.clone.1, %constant.1012), dimensions={0,1,2}, to_apply=%region_23.28, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.155 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.125, %reduce.126.clone.1) } -%fused_computation.244 (param_0.694: f32[14336,4,4096]) -> bf16[4,14336,4096] { - %param_0.694 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.694), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.323 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.245 (param_0.681: f32[14336,4,4096]) -> bf16[4,14336,4096] { + %param_0.681 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) + %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.681), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} + ROOT %bitcast.317 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.245 (param_0.696: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.696 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.696), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.324 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.246 (param_0.683: f32[4096,4,14336]) -> bf16[4,4096,14336] { + %param_0.683 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.683), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} + ROOT %bitcast.318 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.246 (param_0.698: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.698 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.698), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.325 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.247 (param_0.685: f32[4096,4,14336]) -> bf16[4,4096,14336] { + %param_0.685 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.685), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} + ROOT %bitcast.319 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_52.57 (reduce_sum.289: f32[], reduce_sum.290: f32[]) -> f32[] { @@ -312,39 +312,39 @@ StackFrames ROOT %reduce_sum.219 = f32[]{:T(128)} add(%reduce_sum.217, %reduce_sum.218), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.247 (param_0.1114: f32[14336,4,4096], param_1.1277: f32[], param_2.1108: f32[], param_3.797: f32[], param_4.498: f32[14336,4,4096], param_5.423: f32[], param_6.295: f32[4,14336,4096], param_7.194: pred[], param_8.112: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { - %param_0.1114 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %param_3.797 = f32[]{:T(128)S(6)} parameter(3) - %mul.1550.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.797), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.194 = pred[]{:T(512)S(6)} parameter(7) - %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.194), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.295 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) - %bitcast.423.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.295), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.423 = f32[]{:T(128)} parameter(5) - %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.423), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.423.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.423.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.248 (param_0.1100: f32[14336,4,4096], param_1.1263: f32[], param_2.1099: f32[], param_3.793: f32[], param_4.493: f32[14336,4,4096], param_5.415: f32[], param_6.283: f32[4,14336,4096], param_7.182: pred[], param_8.108: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { + %param_0.1100 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) + %param_3.793 = f32[]{:T(128)S(6)} parameter(3) + %mul.1544.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.182 = pred[]{:T(512)S(6)} parameter(7) + %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.182), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.283 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) + %bitcast.417.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.415 = f32[]{:T(128)} parameter(5) + %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.417.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.417.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.955.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.586.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.955.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1556.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.112 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) + %mul.1550.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.108 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) %constant.959.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1557.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1555.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.112, %mul.1557.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1556.clone.1, %mul.1555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1108 = f32[]{:T(128)S(6)} parameter(2) - %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1108), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1551.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1549.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.108, %mul.1551.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1550.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1099 = f32[]{:T(128)S(6)} parameter(2) + %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1099), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.68.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %select_n.273.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.958.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1554.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1552.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.498 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) + %mul.1548.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1546.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1548.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.493 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) %constant.957.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1553.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1551.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.498, %mul.1553.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1552.clone.1, %mul.1551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1277 = f32[]{:T(128)S(6)} parameter(1) - %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1547.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1545.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.493, %mul.1547.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1546.clone.1, %mul.1545.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1263 = f32[]{:T(128)S(6)} parameter(1) + %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1263), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.783.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.819.clone.1, %div.784.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.66.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} sqrt(%div.783.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.956.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -352,10 +352,10 @@ StackFrames %add.817.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%sqrt.66.clone.1, %add.818.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.264.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%div.785.clone.1, %add.817.clone.1), metadata={op_name="multiply.34"} %div.782.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.820.clone.1, %multiply.264.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1549.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1114, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1548.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1550.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1114, %mul.1548.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1543.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1100, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1543.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1542.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1544.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1100, %mul.1542.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.171 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%add.815.clone.1, %add.815.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1001 = f32[]{:T(128)} constant(0) %reduce.127 = f32[]{:T(128)} reduce(%square.171, %constant.1001), dimensions={0,1,2}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -375,39 +375,39 @@ StackFrames ROOT %reduce_sum.213 = f32[]{:T(128)} add(%reduce_sum.211, %reduce_sum.212), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.248 (param_0.1115: f32[4096,4,14336], param_1.1278: f32[], param_2.1109: f32[], param_3.798: f32[], param_4.499: f32[4096,4,14336], param_5.424: f32[], param_6.296: f32[4,4096,14336], param_7.195: pred[], param_8.113: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1115 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.798 = f32[]{:T(128)S(6)} parameter(3) - %mul.1560.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.798), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.195 = pred[]{:T(512)S(6)} parameter(7) - %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.195), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.296 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.425.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.296), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.424 = f32[]{:T(128)} parameter(5) - %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.424), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.425.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.425.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.249 (param_0.1101: f32[4096,4,14336], param_1.1264: f32[], param_2.1100: f32[], param_3.794: f32[], param_4.494: f32[4096,4,14336], param_5.416: f32[], param_6.284: f32[4,4096,14336], param_7.183: pred[], param_8.109: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { + %param_0.1101 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.794 = f32[]{:T(128)S(6)} parameter(3) + %mul.1554.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.183 = pred[]{:T(512)S(6)} parameter(7) + %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.183), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.419.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.416 = f32[]{:T(128)} parameter(5) + %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.419.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.419.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.961.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.592.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.961.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1564.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.113 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) + %mul.1558.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.109 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) %constant.965.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.591.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.965.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1563.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.113, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1564.clone.1, %mul.1563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1109 = f32[]{:T(128)S(6)} parameter(2) - %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1109), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1557.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.109, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1558.clone.1, %mul.1557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) + %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.69.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %select_n.277.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.964.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.590.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.964.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1562.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.499 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) + %mul.1556.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.494 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) %constant.963.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.589.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.963.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1561.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.499, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1562.clone.1, %mul.1561.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1278 = f32[]{:T(128)S(6)} parameter(1) - %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1278), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1555.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.494, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1556.clone.1, %mul.1555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1264 = f32[]{:T(128)S(6)} parameter(1) + %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1264), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.791.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.824.clone.1, %div.792.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.67.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.791.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.962.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -415,10 +415,10 @@ StackFrames %add.823.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.67.clone.1, %broadcast.587.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.265.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.793.clone.1, %add.823.clone.1), metadata={op_name="multiply.33"} %div.790.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.825.clone.1, %multiply.265.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1559.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1115, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1558.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1560.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1115, %mul.1558.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1553.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1101, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1553.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1552.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1554.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1101, %mul.1552.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.172 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.821.clone.1, %add.821.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1002 = f32[]{:T(128)} constant(0) %reduce.128 = f32[]{:T(128)} reduce(%square.172, %constant.1002), dimensions={0,1,2}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -438,39 +438,39 @@ StackFrames ROOT %reduce_sum.210 = f32[]{:T(128)} add(%reduce_sum.205, %reduce_sum.206), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.249 (param_0.1116: f32[4096,4,14336], param_1.1279: f32[], param_2.1110: f32[], param_3.799: f32[], param_4.500: f32[4096,4,14336], param_5.425: f32[], param_6.297: f32[4,4096,14336], param_7.196: pred[], param_8.114: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1116 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.799 = f32[]{:T(128)S(6)} parameter(3) - %mul.1567.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.799), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.196 = pred[]{:T(512)S(6)} parameter(7) - %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.196), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.297 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.427.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.297), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.425 = f32[]{:T(128)} parameter(5) - %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.425), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.427.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.427.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.250 (param_0.1102: f32[4096,4,14336], param_1.1265: f32[], param_2.1101: f32[], param_3.795: f32[], param_4.495: f32[4096,4,14336], param_5.417: f32[], param_6.285: f32[4,4096,14336], param_7.184: pred[], param_8.110: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { + %param_0.1102 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.795 = f32[]{:T(128)S(6)} parameter(3) + %mul.1561.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.184 = pred[]{:T(512)S(6)} parameter(7) + %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.184), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.285 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.421.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.417 = f32[]{:T(128)} parameter(5) + %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.421.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.421.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.967.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.598.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.967.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1571.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.114 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) + %mul.1565.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.110 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) %constant.971.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.597.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.971.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1570.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.114, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1571.clone.1, %mul.1570.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1110 = f32[]{:T(128)S(6)} parameter(2) - %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1110), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1564.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.110, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1565.clone.1, %mul.1564.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) + %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.70.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %select_n.281.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.970.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.596.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.970.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1569.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.500 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) + %mul.1563.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.495 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) %constant.969.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.595.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.969.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1568.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.500, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1569.clone.1, %mul.1568.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1279 = f32[]{:T(128)S(6)} parameter(1) - %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1279), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1562.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.495, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1563.clone.1, %mul.1562.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1265 = f32[]{:T(128)S(6)} parameter(1) + %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1265), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.799.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.829.clone.1, %div.800.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.68.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.799.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.968.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -478,10 +478,10 @@ StackFrames %add.828.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.68.clone.1, %broadcast.593.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.266.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.801.clone.1, %add.828.clone.1), metadata={op_name="multiply.32"} %div.798.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.830.clone.1, %multiply.266.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1566.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1116, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1565.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1567.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1116, %mul.1565.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1560.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1102, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1560.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1559.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1561.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1102, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.173 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.826.clone.1, %add.826.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1003 = f32[]{:T(128)} constant(0) %reduce.129 = f32[]{:T(128)} reduce(%square.173, %constant.1003), dimensions={0,1,2}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -495,10 +495,10 @@ StackFrames ROOT %reduce_sum.183 = f32[]{:T(128)} add(%reduce_sum.178, %reduce_sum.182), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.267 (param_0.1120: f32[4,4096,32,128]) -> f32[] { - %param_0.1120 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.329 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1120), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.176 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.329, %bitcast.329), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.268 (param_0.1106: f32[4,4096,32,128]) -> f32[] { + %param_0.1106 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.323 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.176 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.323, %bitcast.323), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1007 = f32[]{:T(128)} constant(0) ROOT %reduce.133 = f32[]{:T(128)} reduce(%square.176, %constant.1007), dimensions={0,1,2,3}, to_apply=%region_30.35, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } @@ -509,18 +509,18 @@ StackFrames ROOT %reduce_sum.177 = f32[]{:T(128)} add(%reduce_sum.175, %reduce_sum.176), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.269 (param_0.1121: f32[4,32,128,4096]) -> f32[] { - %param_0.1121 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.333 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1121), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.179 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.333, %bitcast.333), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.270 (param_0.1107: f32[4,32,128,4096]) -> f32[] { + %param_0.1107 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.327 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.179 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.327, %bitcast.327), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1008 = f32[]{:T(128)} constant(0) ROOT %reduce.134 = f32[]{:T(128)} reduce(%square.179, %constant.1008), dimensions={0,1,2,3}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.270 (param_0.748: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { - %param_0.748 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.748), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.334 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.271 (param_0.735: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { + %param_0.735 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) + %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.735), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} + ROOT %bitcast.328 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_57.62 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { @@ -535,39 +535,39 @@ StackFrames ROOT %reduce_sum.246 = f32[]{:T(128)} add(%reduce_sum.241, %reduce_sum.245), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.271 (param_0.1109: f32[4096,4,32,128], param_1.1272: f32[], param_2.1103: f32[], param_3.792: f32[], param_4.493: f32[4096,4,32,128], param_5.418: f32[], param_6.290: f32[4,4096,32,128], param_7.189: pred[], param_8.107: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { - %param_0.1109 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.792 = f32[]{:T(128)S(6)} parameter(3) - %mul.1509.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.189 = pred[]{:T(512)S(6)} parameter(7) - %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.189), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.290 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.413.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.290), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.418 = f32[]{:T(128)} parameter(5) - %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.413.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.413.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.272 (param_0.1095: f32[4096,4,32,128], param_1.1258: f32[], param_2.1094: f32[], param_3.788: f32[], param_4.488: f32[4096,4,32,128], param_5.410: f32[], param_6.278: f32[4,4096,32,128], param_7.177: pred[], param_8.103: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { + %param_0.1095 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.788 = f32[]{:T(128)S(6)} parameter(3) + %mul.1503.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.788), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.177 = pred[]{:T(512)S(6)} parameter(7) + %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.177), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.278 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.407.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.278), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.410 = f32[]{:T(128)} parameter(5) + %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.410), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.407.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.407.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.925.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.564.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.925.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %mul.1515.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.107 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1509.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.103 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) %constant.929.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1516.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1514.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.107, %mul.1516.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1515.clone.1, %mul.1514.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1103 = f32[]{:T(128)S(6)} parameter(2) - %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1103), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1510.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1508.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.103, %mul.1510.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1509.clone.1, %mul.1508.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1094 = f32[]{:T(128)S(6)} parameter(2) + %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1094), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.63.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %select_n.253.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.928.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1513.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1511.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1513.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.493 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1507.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1505.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1507.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.488 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) %constant.927.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1512.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1510.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.493, %mul.1512.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1511.clone.1, %mul.1510.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) - %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1506.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1504.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.488, %mul.1506.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1505.clone.1, %mul.1504.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1258 = f32[]{:T(128)S(6)} parameter(1) + %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1258), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.743.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.792.clone.1, %div.744.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.61.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} sqrt(%div.743.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.926.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -575,10 +575,10 @@ StackFrames %add.790.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%sqrt.61.clone.1, %add.791.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.259.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%div.745.clone.1, %add.790.clone.1), metadata={op_name="multiply.39"} %div.742.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.793.clone.1, %multiply.259.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1508.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1109, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1508.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1507.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1509.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1109, %mul.1507.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1502.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1095, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1502.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1501.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1503.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1095, %mul.1501.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.180 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%add.788.clone.1, %add.788.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.996 = f32[]{:T(128)} constant(0) %reduce.135 = f32[]{:T(128)} reduce(%square.180, %constant.996), dimensions={0,1,2,3}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -598,39 +598,39 @@ StackFrames ROOT %reduce_sum.240 = f32[]{:T(128)} add(%reduce_sum.238, %reduce_sum.239), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.272 (param_0.1110: f32[32,4,128,4096], param_1.1273: f32[], param_2.1104: f32[], param_3.793: f32[], param_4.494: f32[32,4,128,4096], param_5.419: f32[], param_6.291: f32[4,32,128,4096], param_7.190: pred[], param_8.108: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { - %param_0.1110 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %param_3.793 = f32[]{:T(128)S(6)} parameter(3) - %mul.1519.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.190 = pred[]{:T(512)S(6)} parameter(7) - %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.190), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.291 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.415.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.291), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.419 = f32[]{:T(128)} parameter(5) - %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.419), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.415.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.415.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.273 (param_0.1096: f32[32,4,128,4096], param_1.1259: f32[], param_2.1095: f32[], param_3.789: f32[], param_4.489: f32[32,4,128,4096], param_5.411: f32[], param_6.279: f32[4,32,128,4096], param_7.178: pred[], param_8.104: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { + %param_0.1096 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) + %param_3.789 = f32[]{:T(128)S(6)} parameter(3) + %mul.1513.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.178 = pred[]{:T(512)S(6)} parameter(7) + %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.178), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.279 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.409.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.411 = f32[]{:T(128)} parameter(5) + %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.411), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.409.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.409.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.931.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.566.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.931.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %mul.1525.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.108 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) + %mul.1519.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.104 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) %constant.935.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1526.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1524.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.108, %mul.1526.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1525.clone.1, %mul.1524.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1104 = f32[]{:T(128)S(6)} parameter(2) - %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1104), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1520.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1518.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.104, %mul.1520.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1519.clone.1, %mul.1518.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1095 = f32[]{:T(128)S(6)} parameter(2) + %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1095), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.64.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %select_n.257.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.934.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1523.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1521.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1523.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.494 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) + %mul.1517.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1515.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1517.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.489 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) %constant.933.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1522.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1520.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.494, %mul.1522.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1521.clone.1, %mul.1520.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1273 = f32[]{:T(128)S(6)} parameter(1) - %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1273), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1516.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1514.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.489, %mul.1516.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1515.clone.1, %mul.1514.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1259 = f32[]{:T(128)S(6)} parameter(1) + %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1259), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.751.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.798.clone.1, %div.752.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.62.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} sqrt(%div.751.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.932.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -638,10 +638,10 @@ StackFrames %add.796.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%sqrt.62.clone.1, %add.797.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.260.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%div.753.clone.1, %add.796.clone.1), metadata={op_name="multiply.38"} %div.750.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.799.clone.1, %multiply.260.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1518.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1110, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1518.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1517.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1519.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1110, %mul.1517.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1512.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1096, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1512.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1511.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1513.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1096, %mul.1511.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.181 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%add.794.clone.1, %add.794.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.997 = f32[]{:T(128)} constant(0) %reduce.136 = f32[]{:T(128)} reduce(%square.181, %constant.997), dimensions={0,1,2,3}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -655,23 +655,23 @@ StackFrames ROOT %reduce_sum.267 = f32[]{:T(128)} add(%reduce_sum.262, %reduce_sum.266), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.279 (param_0.1129: bf16[4,128,128256], param_1.1288: f32[4,128], param_2.1115: s32[4,128], param_3.803: bf16[4,128]) -> f32[4,128] { - %param_2.1115 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1115), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.280 (param_0.1115: bf16[4,128,128256], param_1.1274: f32[4,128], param_2.1106: s32[4,128], param_3.799: bf16[4,128]) -> f32[4,128] { + %param_2.1106 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1106), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.25 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.24 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1129 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.950 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1129), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.803 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.803), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.950, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_0.1115 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.938 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1115), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_3.799 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) + %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.799), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.938, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1274 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.60 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %constant.1017 = f32[]{:T(128)} constant(0) %broadcast.511 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%constant.1017), dimensions={}, metadata={op_name="broadcast.83"} - %mul.1373 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1373, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1367 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1367, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_7.10 (reduce_sum.93: f32[], reduce_sum.94: f32[]) -> f32[] { @@ -680,12 +680,12 @@ StackFrames ROOT %reduce_sum.95 = f32[]{:T(128)} add(%reduce_sum.93, %reduce_sum.94), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.284 (param_0.1130: bf16[4,128,128256], param_1.1289: bf16[4,128]) -> f32[4,128] { - %param_0.1130 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.956 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1130), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1289 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1289), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.956, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.285 (param_0.1116: bf16[4,128,128256], param_1.1275: bf16[4,128]) -> f32[4,128] { + %param_0.1116 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.944 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1116), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1275 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1275), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.944, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.54 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} %constant.1018 = f32[]{:T(128)} constant(0) ROOT %reduce.138 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1018), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} @@ -703,23 +703,23 @@ StackFrames ROOT %reduce_sum.171 = f32[]{:T(128)} add(%reduce_sum.169, %reduce_sum.170), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.290 (param_0.1122: f32[4,4096,8,128], param_1.1282: f32[4,4096,8,128]) -> (f32[], f32[]) { - %param_0.1122 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.350 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1122), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.184 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.350, %bitcast.350), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.291 (param_0.1108: f32[4,4096,8,128], param_1.1268: f32[4,4096,8,128]) -> (f32[], f32[]) { + %param_0.1108 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.344 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.184 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.344, %bitcast.344), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1009 = f32[]{:T(128)} constant(0) %reduce.141 = f32[]{:T(128)} reduce(%square.184, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1282 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(1) - %bitcast.354.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.187.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.354.clone.1, %bitcast.354.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %param_1.1268 = f32[4,4096,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) + %bitcast.348.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.187.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.348.clone.1, %bitcast.348.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %reduce.142.clone.1 = f32[]{:T(128)} reduce(%square.187.clone.1, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.156 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.141, %reduce.142.clone.1) } -%fused_computation.293 (param_0.807: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.807 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.807), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.355 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.294 (param_0.794: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { + %param_0.794 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.794), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} + ROOT %bitcast.349 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_58.63 (reduce_sum.324: f32[], reduce_sum.325: f32[]) -> f32[] { @@ -734,39 +734,39 @@ StackFrames ROOT %reduce_sum.252 = f32[]{:T(128)} add(%reduce_sum.247, %reduce_sum.248), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.294 (param_0.1108: f32[4096,4,8,128], param_1.1271: f32[], param_2.1102: f32[], param_3.791: f32[], param_4.492: f32[4096,4,8,128], param_5.417: f32[], param_6.289: f32[4,4096,8,128], param_7.188: pred[], param_8.106: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1108 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.791 = f32[]{:T(128)S(6)} parameter(3) - %mul.1502.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.188 = pred[]{:T(512)S(6)} parameter(7) - %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.188), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.289 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.289), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.417 = f32[]{:T(128)} parameter(5) - %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.411.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.295 (param_0.1094: f32[4096,4,8,128], param_1.1257: f32[], param_2.1093: f32[], param_3.787: f32[], param_4.487: f32[4096,4,8,128], param_5.409: f32[], param_6.277: f32[4,4096,8,128], param_7.176: pred[], param_8.102: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { + %param_0.1094 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %param_3.787 = f32[]{:T(128)S(6)} parameter(3) + %mul.1496.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.787), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.176 = pred[]{:T(512)S(6)} parameter(7) + %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.176), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.277 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.405.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.277), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.409 = f32[]{:T(128)} parameter(5) + %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.405.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.405.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.919.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.562.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.919.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1506.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.106 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1500.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.923.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.561.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.923.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1505.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.106, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1506.clone.1, %mul.1505.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) - %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1499.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.102, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1500.clone.1, %mul.1499.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1093 = f32[]{:T(128)S(6)} parameter(2) + %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1093), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.62.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %select_n.249.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.922.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.560.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.922.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1504.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.492 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1498.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.487 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.921.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.559.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.921.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1503.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.492, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1504.clone.1, %mul.1503.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1271 = f32[]{:T(128)S(6)} parameter(1) - %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1271), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1497.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.487, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1498.clone.1, %mul.1497.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1257 = f32[]{:T(128)S(6)} parameter(1) + %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1257), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.735.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.786.clone.1, %div.736.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.60.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.735.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.920.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -774,15 +774,15 @@ StackFrames %add.785.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.60.clone.1, %broadcast.557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.258.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.737.clone.1, %add.785.clone.1), metadata={op_name="multiply.40"} %div.734.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.787.clone.1, %multiply.258.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1501.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1108, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1501.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1500.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1502.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1108, %mul.1500.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1495.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1094, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1495.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1494.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1496.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1094, %mul.1494.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.188 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.783.clone.1, %add.783.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.995 = f32[]{:T(128)} constant(0) %reduce.143 = f32[]{:T(128)} reduce(%square.188, %constant.995), dimensions={0,1,2,3}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.145.clone.1 = f32[]{:T(128)} reduce(%integer_pow.62.clone.1, %constant.995), dimensions={0,1,2,3}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) + ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) } %region_55.60 (reduce_sum.304: f32[], reduce_sum.308: f32[]) -> f32[] { @@ -797,39 +797,39 @@ StackFrames ROOT %reduce_sum.234 = f32[]{:T(128)} add(%reduce_sum.232, %reduce_sum.233), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.295 (param_0.1111: f32[4096,4,8,128], param_1.1274: f32[], param_2.1105: f32[], param_3.794: f32[], param_4.495: f32[4096,4,8,128], param_5.420: f32[], param_6.292: f32[4,4096,8,128], param_7.191: pred[], param_8.109: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1111 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.794 = f32[]{:T(128)S(6)} parameter(3) - %mul.1529.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.191 = pred[]{:T(512)S(6)} parameter(7) - %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.191), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.292 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.417.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.292), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.420 = f32[]{:T(128)} parameter(5) - %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.420), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.417.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.417.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.296 (param_0.1097: f32[4096,4,8,128], param_1.1260: f32[], param_2.1096: f32[], param_3.790: f32[], param_4.490: f32[4096,4,8,128], param_5.412: f32[], param_6.280: f32[4,4096,8,128], param_7.179: pred[], param_8.105: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { + %param_0.1097 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.790 = f32[]{:T(128)S(6)} parameter(3) + %mul.1523.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.179 = pred[]{:T(512)S(6)} parameter(7) + %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.179), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.280 = f32[4,4096,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.412 = f32[]{:T(128)} parameter(5) + %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.411.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.937.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.572.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.937.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1533.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.109 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1527.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.105 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.941.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.571.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.941.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1532.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.109, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1533.clone.1, %mul.1532.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1105 = f32[]{:T(128)S(6)} parameter(2) - %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1105), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1526.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.105, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1527.clone.1, %mul.1526.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1096 = f32[]{:T(128)S(6)} parameter(2) + %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1096), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.65.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %select_n.261.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.940.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.570.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.940.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1531.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.495 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1525.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.490 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.939.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.569.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.939.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1530.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.495, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1531.clone.1, %mul.1530.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1274 = f32[]{:T(128)S(6)} parameter(1) - %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1524.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.490, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1525.clone.1, %mul.1524.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1260 = f32[]{:T(128)S(6)} parameter(1) + %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1260), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.759.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.803.clone.1, %div.760.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.63.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.759.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.938.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -837,10 +837,10 @@ StackFrames %add.802.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.761.clone.1, %add.802.clone.1), metadata={op_name="multiply.37"} %div.758.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.804.clone.1, %multiply.261.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1528.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1111, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1528.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1527.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1529.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1111, %mul.1527.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1522.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1097, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1522.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1521.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1523.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1097, %mul.1521.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.189 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.800.clone.1, %add.800.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.998 = f32[]{:T(128)} constant(0) %reduce.144 = f32[]{:T(128)} reduce(%square.189, %constant.998), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -848,22 +848,22 @@ StackFrames ROOT %tuple.143 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.144, %add.800.clone.1, %add.803.clone.1, %add.804.clone.1, %reduce.146.clone.1) } -%fused_computation.311 (param_0.872: bf16[4,128,4096], param_1.941: f32[4,128], param_2.726: f32[4,128], param_3.452: bf16[4,128,4096], param_4.271: bf16[4096]) -> bf16[4,128,4096] { - %param_3.452 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.271 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.375 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.271), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.365 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.452, %dot_general.375), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.973 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.726 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1423 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.726), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1415 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.973, %mul.1423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.872 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.984 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.872), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.941 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1422 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.941), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1421 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.984, %mul.1422), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1415, %mul.1421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} - ROOT %convert_element_type.971 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.312 (param_0.859: bf16[4,128,4096], param_1.928: f32[4,128], param_2.717: f32[4,128], param_3.448: bf16[4,128,4096], param_4.266: bf16[4096]) -> bf16[4,128,4096] { + %param_3.448 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.266 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %dot_general.371 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.266), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.361 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.448, %dot_general.371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.961 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_2.717 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %mul.1417 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.717), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1409 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.961, %mul.1417), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_0.859 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.972 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.859), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.928 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1416 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.928), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1415 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.972, %mul.1416), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1409, %mul.1415), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} + ROOT %convert_element_type.959 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} } %region_5.8 (reduce_sum.87: f32[], reduce_sum.88: f32[]) -> f32[] { @@ -872,10 +872,10 @@ StackFrames ROOT %reduce_sum.92 = f32[]{:T(128)} add(%reduce_sum.87, %reduce_sum.88), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.312 (param_0.1131: bf16[4,128,4096]) -> f32[4,128] { - %param_0.1131 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.975 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1131), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.192 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.975, %convert_element_type.975), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} +%fused_computation.313 (param_0.1117: bf16[4,128,4096]) -> f32[4,128] { + %param_0.1117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.963 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1117), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.192 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.963, %convert_element_type.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} %constant.1019 = f32[]{:T(128)} constant(0) ROOT %reduce.147 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.192, %constant.1019), dimensions={2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } @@ -886,17 +886,17 @@ StackFrames ROOT %reduce_sum.107 = f32[]{:T(128)} add(%reduce_sum.102, %reduce_sum.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.314 (param_0.1126: bf16[4,128,4096], param_1.1285: bf16[4,128,4096], param_2.1113: bf16[4096]) -> f32[4,128] { - %param_0.1126 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1126), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1285 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.374 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1113), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.364 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1285, %dot_general.374), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.981 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1419 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %convert_element_type.981), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} +%fused_computation.315 (param_0.1112: bf16[4,128,4096], param_1.1271: bf16[4,128,4096], param_2.1104: bf16[4096]) -> f32[4,128] { + %param_0.1112 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.970 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1112), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1104 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.370 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1104), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.360 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1271, %dot_general.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.969 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.1413 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.970, %convert_element_type.969), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.1013 = f32[]{:T(128)} constant(0) - ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1419, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1413, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_8.11 (dot_general.182: bf16[], dot_general.183: bf16[]) -> bf16[] { @@ -905,86 +905,86 @@ StackFrames ROOT %add.168 = bf16[]{:T(256)} add(%dot_general.182, %dot_general.183), metadata={op_name="add.54"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.235.clone.clone (param_0.1095: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1095 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1033 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1095), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.449 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +%fused_computation.236.clone.clone (param_0.1081: f32[4096,128256]) -> bf16[4096,128256,1] { + %param_0.1081 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %convert_element_type.1021 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1081), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + ROOT %bitcast.443 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1021), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} } -%fused_computation.280.clone.1.clone.clone (param_0.1096: bf16[4,128,128256], param_1.1261: s32[4,128], param_2.1081: f32[4,128], param_3.782: f32[4,128], param_4.484: bf16[4,128], param_5.409: f32[4,128]) -> bf16[4,128,128256] { - %param_5.409 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1603 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.782 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1602 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.782), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1096 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1036 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1096), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.484 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.484), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1036, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.281.clone.1.clone.clone (param_0.1082: bf16[4,128,128256], param_1.1247: s32[4,128], param_2.1072: f32[4,128], param_3.778: f32[4,128], param_4.479: bf16[4,128], param_5.401: f32[4,128]) -> bf16[4,128,128256] { + %param_5.401 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.1597 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.401), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.778 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.1596 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.778), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1082 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1024 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.479 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.479), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1024, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.60 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1601 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1602, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1081), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1601, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1261 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1261), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.1595 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1596, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1072 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1072), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1595, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1247 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1247), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.42 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.41 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1035 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1035), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1600 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1603, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1034 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1600), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.315 (param_0.1094: f32[4,128], param_1.1260: bf16[4,128,4096], param_2.1082: f32[4096,128256], param_3.783: bf16[4,128,128256], param_4.485: s32[4,128], param_5.410: f32[4,128], param_6.284: f32[4,128], param_7.183: bf16[4,128], param_8.102: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { - %param_3.783 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.485 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.410 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.284 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.183 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.102 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.783, %param_4.485, %param_5.410, %param_6.284, %param_7.183, /*index=5*/%param_8.102), kind=kLoop, calls=%fused_computation.280.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1082 = f32[4096,128256]{1,0:T(8,128)} parameter(2) - %fusion.219.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1082), kind=kLoop, calls=%fused_computation.235.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.86.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.219.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %param_1.1260 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.994 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1260), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1094 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1434 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1094), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1433 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.994, %mul.1434), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.993 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1433), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.86.clone.1, %convert_element_type.993), metadata={op_name="multiply.206"} + %convert_element_type.1023 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1023), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1594 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1597, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1022 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1594), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.316 (param_0.1080: f32[4,128], param_1.1246: bf16[4,128,4096], param_2.1073: f32[4096,128256], param_3.779: bf16[4,128,128256], param_4.480: s32[4,128], param_5.402: f32[4,128], param_6.272: f32[4,128], param_7.171: bf16[4,128], param_8.98: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { + %param_3.779 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) + %param_4.480 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %param_5.402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.272 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.171 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) + %param_8.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) + %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.779, %param_4.480, %param_5.402, %param_6.272, %param_7.171, /*index=5*/%param_8.98), kind=kLoop, calls=%fused_computation.281.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_2.1073 = f32[4096,128256]{1,0:T(8,128)} parameter(2) + %fusion.209.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1073), kind=kLoop, calls=%fused_computation.236.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %convolution.80.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.209.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %param_1.1246 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1246), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1080 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1428 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1080), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1427 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %mul.1428), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.981 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1427), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.80.clone.1, %convert_element_type.981), metadata={op_name="multiply.206"} %constant.874 = bf16[]{:T(256)} constant(0) %reduce.149 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.252, %constant.874), dimensions={0,1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.86.clone.1) + ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.80.clone.1) } -%fused_computation.323 (param_0.904: f32[64], param_1.974: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.974), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.904 = f32[64]{0:T(128)S(1)} parameter(0) - %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.904), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} +%fused_computation.324 (param_0.891: f32[64], param_1.961: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.961 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.961), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %param_0.891 = f32[64]{0:T(128)S(1)} parameter(0) + %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.891), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %div.618 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.621, %div.619), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.618), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} - %convert_element_type.1002 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.990 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %cos.41.clone.1 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.618), metadata={op_name="jit(train_step)/layers/cos" stack_frame_id=0} - %convert_element_type.1001.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1002, %convert_element_type.1001.clone.1) + %convert_element_type.989.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.990, %convert_element_type.989.clone.1) } -%fused_computation.324 (param_0.901: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.901 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.325 (param_0.888: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.888 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.866 = bf16[]{:T(256)} constant(-inf) - %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.888, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.888, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.34 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.38, %pad.37), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } -%fused_computation.325 (param_0.903: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.903 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.326 (param_0.890: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.890 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.865 = bf16[]{:T(256)} constant(-inf) - %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.890, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.890, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.35 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.40, %pad.39), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } @@ -1000,15 +1000,15 @@ StackFrames ROOT %reduce_sum.162 = f32[]{:T(128)} add(%reduce_sum.157, %reduce_sum.161), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.329 (param_0.1123: f32[4,4096], param_1.1283: f32[4,4096]) -> (f32[], f32[]) { - %param_0.1123 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.371 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.195 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.330 (param_0.1109: f32[4,4096], param_1.1269: f32[4,4096]) -> (f32[], f32[]) { + %param_0.1109 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.365 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.195 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.365, %bitcast.365), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1010 = f32[]{:T(128)} constant(0) %reduce.150 = f32[]{:T(128)} reduce(%square.195, %constant.1010), dimensions={0,1}, to_apply=%region_27.32, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1283 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) - %bitcast.375.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.198.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %param_1.1269 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) + %bitcast.369.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.198.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.369.clone.1, %bitcast.369.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %reduce.151.clone.1 = f32[]{:T(128)} reduce(%square.198.clone.1, %constant.1010), dimensions={0,1}, to_apply=%region_26.31, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.150, %reduce.151.clone.1) } @@ -1025,39 +1025,39 @@ StackFrames ROOT %reduce_sum.231 = f32[]{:T(128)} add(%reduce_sum.226, %reduce_sum.227), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.332 (param_0.1112: f32[4096,4], param_1.1275: f32[], param_2.1106: f32[], param_3.795: f32[], param_4.496: f32[4096,4], param_5.421: f32[], param_6.293: f32[4,4096], param_7.192: pred[], param_8.110: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1112 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.795 = f32[]{:T(128)S(6)} parameter(3) - %mul.1536.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.192 = pred[]{:T(512)S(6)} parameter(7) - %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.192), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.293 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.419.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.293), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.421 = f32[]{:T(128)} parameter(5) - %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.421), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.419.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.419.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.333 (param_0.1098: f32[4096,4], param_1.1261: f32[], param_2.1097: f32[], param_3.791: f32[], param_4.491: f32[4096,4], param_5.413: f32[], param_6.281: f32[4,4096], param_7.180: pred[], param_8.106: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1098 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.791 = f32[]{:T(128)S(6)} parameter(3) + %mul.1530.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.180 = pred[]{:T(512)S(6)} parameter(7) + %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.180), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.281 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.413.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.281), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.413 = f32[]{:T(128)} parameter(5) + %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.413), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.413.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.413.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.943.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.578.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.943.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1540.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.110 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1534.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.106 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) %constant.947.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.577.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.947.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1539.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.110, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1540.clone.1, %mul.1539.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1106 = f32[]{:T(128)S(6)} parameter(2) - %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1106), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1533.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.106, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1534.clone.1, %mul.1533.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1097 = f32[]{:T(128)S(6)} parameter(2) + %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1097), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.66.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %select_n.265.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.946.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.576.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.946.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1538.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.496 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1532.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.491 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) %constant.945.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.575.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.945.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1537.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.496, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1538.clone.1, %mul.1537.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1275 = f32[]{:T(128)S(6)} parameter(1) - %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1275), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1531.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.491, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1532.clone.1, %mul.1531.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1261 = f32[]{:T(128)S(6)} parameter(1) + %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1261), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.767.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.808.clone.1, %div.768.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.64.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.767.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.944.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1065,10 +1065,10 @@ StackFrames %add.807.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.262.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.769.clone.1, %add.807.clone.1), metadata={op_name="multiply.36"} %div.766.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.809.clone.1, %multiply.262.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1535.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1112, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1535.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1534.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1536.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1112, %mul.1534.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1529.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1098, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1529.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1528.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1530.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1098, %mul.1528.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.199 = f32[4096,4]{0,1:T(4,128)} multiply(%add.805.clone.1, %add.805.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.999 = f32[]{:T(128)} constant(0) %reduce.152 = f32[]{:T(128)} reduce(%square.199, %constant.999), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1088,39 +1088,39 @@ StackFrames ROOT %reduce_sum.225 = f32[]{:T(128)} add(%reduce_sum.220, %reduce_sum.224), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.333 (param_0.1113: f32[4096,4], param_1.1276: f32[], param_2.1107: f32[], param_3.796: f32[], param_4.497: f32[4096,4], param_5.422: f32[], param_6.294: f32[4,4096], param_7.193: pred[], param_8.111: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1113 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.796 = f32[]{:T(128)S(6)} parameter(3) - %mul.1543.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.193 = pred[]{:T(512)S(6)} parameter(7) - %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.193), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.294 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.421.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.294), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.422 = f32[]{:T(128)} parameter(5) - %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.422), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.421.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.421.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.334 (param_0.1099: f32[4096,4], param_1.1262: f32[], param_2.1098: f32[], param_3.792: f32[], param_4.492: f32[4096,4], param_5.414: f32[], param_6.282: f32[4,4096], param_7.181: pred[], param_8.107: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1099 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.792 = f32[]{:T(128)S(6)} parameter(3) + %mul.1537.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.181 = pred[]{:T(512)S(6)} parameter(7) + %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.181), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.282 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.415.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.414 = f32[]{:T(128)} parameter(5) + %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.414), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.415.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.415.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.949.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.584.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.949.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1547.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.111 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1541.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.107 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) %constant.953.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.583.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.953.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1546.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.111, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1547.clone.1, %mul.1546.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1107 = f32[]{:T(128)S(6)} parameter(2) - %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1107), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1540.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.107, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1541.clone.1, %mul.1540.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1098 = f32[]{:T(128)S(6)} parameter(2) + %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1098), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.67.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %select_n.269.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.952.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.582.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.952.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1545.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.497 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1539.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.492 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) %constant.951.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.581.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.951.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1544.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.497, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1545.clone.1, %mul.1544.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1276 = f32[]{:T(128)S(6)} parameter(1) - %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1276), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1538.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.492, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1539.clone.1, %mul.1538.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1262 = f32[]{:T(128)S(6)} parameter(1) + %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1262), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.775.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.813.clone.1, %div.776.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.65.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.775.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.950.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1128,10 +1128,10 @@ StackFrames %add.812.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.65.clone.1, %broadcast.579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.263.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.777.clone.1, %add.812.clone.1), metadata={op_name="multiply.35"} %div.774.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.814.clone.1, %multiply.263.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1542.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1113, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1542.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1541.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1543.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1113, %mul.1541.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1536.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1099, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1536.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1535.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1537.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1099, %mul.1535.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.200 = f32[4096,4]{0,1:T(4,128)} multiply(%add.810.clone.1, %add.810.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1000 = f32[]{:T(128)} constant(0) %reduce.153 = f32[]{:T(128)} reduce(%square.200, %constant.1000), dimensions={0,1}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1145,10 +1145,10 @@ StackFrames ROOT %reduce_sum.101 = f32[]{:T(128)} add(%reduce_sum.99, %reduce_sum.100), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.344 (param_0.1127: bf16[4096]) -> f32[] { - %param_0.1127 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.1006 = f32[4096]{0:T(1024)} convert(%param_0.1127), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.203 = f32[4096]{0:T(1024)} multiply(%convert_element_type.1006, %convert_element_type.1006), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.345 (param_0.1113: bf16[4096]) -> f32[] { + %param_0.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.994 = f32[4096]{0:T(1024)} convert(%param_0.1113), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.203 = f32[4096]{0:T(1024)} multiply(%convert_element_type.994, %convert_element_type.994), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1014 = f32[]{:T(128)} constant(0) ROOT %reduce.156 = f32[]{:T(128)} reduce(%square.203, %constant.1014), dimensions={0}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } @@ -1165,39 +1165,39 @@ StackFrames ROOT %reduce_sum.204 = f32[]{:T(128)} add(%reduce_sum.199, %reduce_sum.203), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.345 (param_0.1117: f32[4096], param_1.1280: f32[], param_2.1111: f32[], param_3.800: f32[], param_4.501: f32[4096], param_5.426: f32[], param_6.298: bf16[4096], param_7.197: pred[], param_8.115: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { - %param_0.1117 = f32[4096]{0:T(1024)S(1)} parameter(0) - %param_3.800 = f32[]{:T(128)S(6)} parameter(3) - %mul.1574.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.800), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.197 = pred[]{:T(512)S(6)} parameter(7) - %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.298 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1021.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.298), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.426 = f32[]{:T(128)} parameter(5) - %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.426), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1021.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1021.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.346 (param_0.1103: f32[4096], param_1.1266: f32[], param_2.1102: f32[], param_3.796: f32[], param_4.496: f32[4096], param_5.418: f32[], param_6.286: bf16[4096], param_7.185: pred[], param_8.111: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { + %param_0.1103 = f32[4096]{0:T(1024)S(1)} parameter(0) + %param_3.796 = f32[]{:T(128)S(6)} parameter(3) + %mul.1568.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.185 = pred[]{:T(512)S(6)} parameter(7) + %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.185), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.286 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1009.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.286), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_5.418 = f32[]{:T(128)} parameter(5) + %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1009.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1009.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.973.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.600.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.973.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.1580.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.115 = f32[4096]{0:T(1024)S(1)} parameter(8) + %mul.1574.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.111 = f32[4096]{0:T(1024)S(1)} parameter(8) %constant.977.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1581.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1579.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.115, %mul.1581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1580.clone.1, %mul.1579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1111 = f32[]{:T(128)S(6)} parameter(2) - %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1111), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1575.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1573.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.111, %mul.1575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1574.clone.1, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) + %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.71.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %select_n.285.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.976.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1578.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1576.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.501 = f32[4096]{0:T(1024)S(1)} parameter(4) + %mul.1572.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1570.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.496 = f32[4096]{0:T(1024)S(1)} parameter(4) %constant.975.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1577.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1575.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.501, %mul.1577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1576.clone.1, %mul.1575.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1280 = f32[]{:T(128)S(6)} parameter(1) - %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1280), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1571.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1569.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.496, %mul.1571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1570.clone.1, %mul.1569.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1266 = f32[]{:T(128)S(6)} parameter(1) + %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1266), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.807.clone.1 = f32[4096]{0:T(1024)} divide(%add.835.clone.1, %div.808.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.69.clone.1 = f32[4096]{0:T(1024)} sqrt(%div.807.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.974.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1205,10 +1205,10 @@ StackFrames %add.833.clone.1 = f32[4096]{0:T(1024)} add(%sqrt.69.clone.1, %add.834.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.267.clone.1 = f32[4096]{0:T(1024)} multiply(%div.809.clone.1, %add.833.clone.1), metadata={op_name="multiply.31"} %div.806.clone.1 = f32[4096]{0:T(1024)} divide(%add.836.clone.1, %multiply.267.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1573.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1117, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1572.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1574.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1117, %mul.1572.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1567.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1103, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1566.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1568.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1103, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.204 = f32[4096]{0:T(1024)} multiply(%add.831.clone.1, %add.831.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1004 = f32[]{:T(128)} constant(0) %reduce.157 = f32[]{:T(128)} reduce(%square.204, %constant.1004), dimensions={0}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1216,27 +1216,27 @@ StackFrames ROOT %tuple.148 = (f32[]{:T(128)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.157, %add.831.clone.1, %add.835.clone.1, %add.836.clone.1, %reduce.158.clone.1) } -%fused_computation.351 (param_0.964: s32[512]) -> s32[1024] { +%fused_computation.352 (param_0.951: s32[512]) -> s32[1024] { %constant.801 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.539 = s32[1024]{0:T(1024)} broadcast(%constant.801), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.964 = s32[512]{0:T(512)S(1)} parameter(0) + %param_0.951 = s32[512]{0:T(512)S(1)} parameter(0) %constant.802 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.964, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.951, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %constant.800 = s32[] constant(128255), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.538 = s32[1024]{0:T(1024)} broadcast(%constant.800), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.539, %pad.41, %broadcast.538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } -%fused_computation.352 (param_0.963: s32[4,128]) -> s32[512] { - %param_0.963 = s32[4,128]{1,0:T(4,128)} parameter(0) +%fused_computation.353 (param_0.950: s32[4,128]) -> s32[512] { + %param_0.950 = s32[4,128]{1,0:T(4,128)} parameter(0) %constant.888 = s32[]{:T(128)} constant(0) %broadcast.546 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.888), dimensions={}, metadata={op_name="broadcast.81"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.963, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} + %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.950, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} %constant.875 = s32[]{:T(128)} constant(128256) %add.760 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.875), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.963, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} - ROOT %bitcast.376 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.950, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} + %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} + ROOT %bitcast.370 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} } %region_61.66 (reduce_sum.345: f32[], reduce_sum.346: f32[]) -> f32[] { @@ -1251,52 +1251,52 @@ StackFrames ROOT %reduce_sum.273 = f32[]{:T(128)} add(%reduce_sum.268, %reduce_sum.269), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.353 (param_0.1128: bf16[4,128], param_1.1287: f32[4,128], param_2.1114: f32[4,128], param_3.802: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { - %param_3.802 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) +%fused_computation.354 (param_0.1114: bf16[4,128], param_1.1273: f32[4,128], param_2.1105: f32[4,128], param_3.798: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { + %param_3.798 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) %constant.979.clone.1 = s32[]{:T(128)} constant(0) %broadcast.601.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.979.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.802, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1287), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1128 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1128), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.798, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %param_1.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1273), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_0.1114 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1114), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} %add.762 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} %square.207 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %add.762), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} %constant.1016 = f32[]{:T(128)} constant(0) %broadcast.543 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1016), dimensions={}, metadata={op_name="broadcast.32"} - %mul.1473 = f32[4,128]{1,0:T(4,128)} multiply(%square.207, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %mul.1465 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1473, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.159 = f32[]{:T(128)} reduce(%mul.1465, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1114), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} - %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1473), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %mul.1466.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1466.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %mul.1471.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1467 = f32[4,128]{1,0:T(4,128)} multiply(%square.207, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1459 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1467, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.159 = f32[]{:T(128)} reduce(%mul.1459, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_2.1105 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1105), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1467), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1460.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1460.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1465.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %constant.891.clone.1 = f32[]{:T(128)} constant(1) %add.757.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.891.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1471.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1465.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} ROOT %tuple.149 = (f32[]{:T(128)}, f32[]{:T(128)}, pred[4,128]{1,0:T(4,128)(4,1)S(1)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%reduce.159, %reduce.160.clone.1, %ne.6.clone.1, %add.750.clone.1) } -%fused_computation.356 (param_0.987: f32[4,128], param_1.1101: f32[4,128]) -> f32[4,128] { - %param_0.987 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1101 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.357 (param_0.974: f32[4,128], param_1.1088: f32[4,128]) -> f32[4,128] { + %param_0.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1088 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.869 = f32[]{:T(128)} constant(0.000244140625) %broadcast.549 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.869), dimensions={}, metadata={op_name="broadcast.264"} - %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1101, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1088, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.867 = f32[]{:T(128)} constant(1e-05) %add.770 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.867), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.769 = f32[4,128]{1,0:T(4,128)} add(%div.656, %add.770), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %rsqrt.90 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} %div.649 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.90, %add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.864 = f32[]{:T(128)} constant(-0.5) - %mul.1477 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1470 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1477), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1469 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.987, %mul.1470), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1471 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1464 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1471), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1463 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.974, %mul.1464), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.863 = f32[]{:T(128)} constant(0.00048828125) - %mul.1476 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - ROOT %mul.1468 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1469, %mul.1476), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1470 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + ROOT %mul.1462 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1463, %mul.1470), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} } %region_0.1 (reduce_sum.67: s32[], reduce_sum.71: s32[]) -> s32[] { @@ -1305,64 +1305,64 @@ StackFrames ROOT %reduce_sum.72 = s32[]{:T(128)} add(%reduce_sum.67, %reduce_sum.71), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.360 (param_0.1004: pred[4,128]) -> s32[] { - %param_0.1004 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1013 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1004), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} +%fused_computation.361 (param_0.991: pred[4,128]) -> s32[] { + %param_0.991 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %convert_element_type.1001 = s32[4,128]{1,0:T(4,128)} convert(%param_0.991), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} %constant.889 = s32[]{:T(128)} constant(0) - ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1013, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1001, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.361 (param_0.989: f32[4,128]) -> f32[4,128] { - %param_0.989 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.362 (param_0.976: f32[4,128]) -> f32[4,128] { + %param_0.976 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.870 = f32[]{:T(128)} constant(0.000244140625) %broadcast.541 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.870), dimensions={}, metadata={op_name="broadcast.264"} - %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.989, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.976, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.868 = f32[]{:T(128)} constant(1e-05) %add.759 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.868), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.756 = f32[4,128]{1,0:T(4,128)} add(%div.654, %add.759), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} ROOT %rsqrt.88 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.756), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} } -%fused_computation.362 (param_0.990: pred[4,128], param_1.1286: f32[]) -> f32[4,128] { - %param_0.990 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1286 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1286), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} +%fused_computation.363 (param_0.977: pred[4,128], param_1.1272: f32[]) -> f32[4,128] { + %param_0.977 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) + %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} %constant.1015 = f32[]{:T(128)} constant(0) %broadcast.545 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1015), dimensions={}, metadata={op_name="broadcast.32"} - ROOT %mul.1478 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.990, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %mul.1472 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.977, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} } -%fused_computation.364 () -> f32[64] { +%fused_computation.365 () -> f32[64] { %constant.873 = f32[]{:T(128)} constant(500000) %broadcast.552 = f32[64]{0:T(128)} broadcast(%constant.873), dimensions={}, metadata={op_name="broadcast.255"} %iota.46 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/layers/iota" stack_frame_id=0} %constant.872 = s32[]{:T(128)} constant(2) %broadcast.551 = s32[64]{0:T(128)} broadcast(%constant.872), dimensions={}, metadata={op_name="broadcast.256"} - %mul.1479 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} - %convert_element_type.1014 = f32[64]{0:T(128)} convert(%mul.1479), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %mul.1473 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} + %convert_element_type.1002 = f32[64]{0:T(128)} convert(%mul.1473), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %constant.871 = f32[]{:T(128)} constant(0.0078125) %broadcast.550 = f32[64]{0:T(128)} broadcast(%constant.871), dimensions={}, metadata={op_name="broadcast.257"} - %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1014, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1002, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.552, %div.657), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} } -%fused_computation.365 (param_0.1002: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.1002 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1015 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1002), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %bitcast.377 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1015), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.377, %convert_element_type.1015) +%fused_computation.366 (param_0.989: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { + %param_0.989 = s32[4,128]{1,0:T(4,128)} parameter(0) + %convert_element_type.1003 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.989), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %bitcast.371 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.371, %convert_element_type.1003) } -%fused_computation.369 (param_0.1103: f32[4096,4]) -> bf16[4,4096] { - %param_0.1103 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.451 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1103), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.451) +%fused_computation.369 (param_0.1089: f32[4096,4]) -> bf16[4,4096] { + %param_0.1089 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.445 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1089), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.445) } -%fused_computation.370 (param_0.1104: f32[4096,4]) -> bf16[4,4096] { - %param_0.1104 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.452 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1104), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.452) +%fused_computation.370 (param_0.1090: f32[4096,4]) -> bf16[4,4096] { + %param_0.1090 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.446 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.446) } %region_6.9 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { @@ -1371,41 +1371,41 @@ StackFrames ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.237.clone.clone (param_0.1090: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1090 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1026 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.447 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1026), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} -} - -%fused_computation.317.clone.clone (param_0.1091: f32[4,128], param_1.1257: bf16[4,128,4096], param_2.1077: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1077 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1077), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1257 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1028 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1257), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1091 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1595 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1091), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1594 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1028, %mul.1595), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1027 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1594), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1027), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.371 (param_0.1105: f32[4096,128256], param_1.1268: f32[4,128], param_2.1099: bf16[4,128,4096], param_3.788: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { - %param_1.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1099 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.788 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.240.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1268, %param_2.1099, %param_3.788), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1105 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %fusion.221.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1105), kind=kLoop, calls=%fused_computation.237.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.87.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.240.clone.1, %fusion.221.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} +%fused_computation.238.clone.clone (param_0.1076: f32[4096,128256]) -> bf16[4096,128256,1] { + %param_0.1076 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %convert_element_type.1014 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1076), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + ROOT %bitcast.441 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1014), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +} + +%fused_computation.318.clone.clone (param_0.1077: f32[4,128], param_1.1243: bf16[4,128,4096], param_2.1068: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1068 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.379 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1068), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1243 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1016 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1243), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1077 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1589 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1077), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1588 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1016, %mul.1589), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1015 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1588), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.378 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.379, %convert_element_type.1015), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.371 (param_0.1091: f32[4096,128256], param_1.1254: f32[4,128], param_2.1090: bf16[4,128,4096], param_3.784: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { + %param_1.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1090 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.784 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.230.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1254, %param_2.1090, %param_3.784), kind=kLoop, calls=%fused_computation.318.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1091 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %fusion.211.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1091), kind=kLoop, calls=%fused_computation.238.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %convolution.81.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.230.clone.1, %fusion.211.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} %constant.992 = bf16[]{:T(256)} constant(-inf) - %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.87.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} - ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.87.clone.1) + %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.81.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.81.clone.1) } -%fused_computation.372 (param_0.1102: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.1102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %bitcast.450 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1102), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.450) +%fused_computation.372 (param_0.1088: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { + %param_0.1088 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %bitcast.444 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1088), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.444) } %convert_element_type.525.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { @@ -1414,13 +1414,13 @@ StackFrames ROOT %add.624 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.121.clone.clone (param_0.1242: bf16[4,4096], param_1.1376: s32[]) -> bf16[4096] { - %param_0.1242 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.121.clone.clone (param_0.1229: bf16[4,4096], param_1.1363: s32[]) -> bf16[4096] { + %param_0.1229 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1363 = s32[]{:T(128)S(6)} parameter(1) %constant.1116 = s32[]{:T(128)} constant(0) - %dynamic_slice.316 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1242, %param_1.1376, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %dynamic_slice.310 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1229, %param_1.1363, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} %constant.1117 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.316, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.310, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %region_12.14 (reduce_sum.108: f32[], reduce_sum.109: f32[]) -> f32[] { @@ -1429,70 +1429,70 @@ StackFrames ROOT %reduce_sum.113 = f32[]{:T(128)} add(%reduce_sum.108, %reduce_sum.109), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.58.clone.clone (param_0.1243: bf16[4,4,128,4096], param_1.1377: s32[]) -> f32[4,128] { - %param_0.1243 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1377 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.58.clone.clone (param_0.1230: bf16[4,4,128,4096], param_1.1364: s32[]) -> f32[4,128] { + %param_0.1230 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1364 = s32[]{:T(128)S(6)} parameter(1) %constant.1118 = s32[]{:T(128)} constant(0) - %dynamic_slice.317 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1243, %param_1.1377, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.548 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1093 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.548), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.214 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1093, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %dynamic_slice.311 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1230, %param_1.1364, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.543 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.311), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1081 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.543), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.214 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1081, %convert_element_type.1081), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} %constant.1119 = f32[]{:T(128)} constant(0) ROOT %reduce.175 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.214, %constant.1119), dimensions={2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} } -%fused_computation.143.clone.1.clone (param_0.1244: f32[4,128]) -> f32[4,128] { - %param_0.1244 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.143.clone.1.clone (param_0.1231: f32[4,128]) -> f32[4,128] { + %param_0.1231 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1121 = f32[]{:T(128)} constant(0.000244140625) %closed_call.81 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1121), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1244, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1231, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1120 = f32[]{:T(128)} constant(1e-05) %closed_call.80 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1120), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.858 = f32[4,128]{1,0:T(4,128)} add(%div.842, %closed_call.80), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.97 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.858), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.24.clone.1.clone.clone (param_0.1258: bf16[4,4096,32,128], param_1.1387: s32[]) -> bf16[4096,32,128,1] { - %param_0.1258 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.24.clone.1.clone.clone (param_0.1245: bf16[4,4096,32,128], param_1.1374: s32[]) -> bf16[4096,32,128,1] { + %param_0.1245 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1374 = s32[]{:T(128)S(6)} parameter(1) %constant.1134 = s32[]{:T(128)} constant(0) - %dynamic_slice.323 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1258, %param_1.1387, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.559 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.317 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1245, %param_1.1374, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.554 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.91.clone.clone (param_0.1259: f32[4,128], param_1.1388: bf16[4,4,128,4096], param_2.1176: s32[], param_3.847: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.847 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.847), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1388 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1176 = s32[]{:T(128)S(6)} parameter(2) +%fused_computation.91.clone.clone (param_0.1246: f32[4,128], param_1.1375: bf16[4,4,128,4096], param_2.1167: s32[], param_3.843: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.843 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.843), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1375 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1167 = s32[]{:T(128)S(6)} parameter(2) %constant.1135 = s32[]{:T(128)} constant(0) - %dynamic_slice.324 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1388, %param_2.1176, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.561 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1101 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1259 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1709 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1259), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1708 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1101, %mul.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1100 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1708), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1100), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.560 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.427), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.36.clone.clone (param_0.1260: bf16[4,4096,32,128], param_1.1389: s32[], param_2.1177: f32[4,128], param_3.848: bf16[4,4,128,4096], param_4.530: bf16[4096]) -> bf16[4,128,32,128] { - %param_2.1177 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.848 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1389 = s32[]{:T(128)S(6)} parameter(1) - %param_4.530 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.343 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1177, %param_3.848, %param_1.1389, %param_4.530), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1260 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.342 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1260, %param_1.1389), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.113 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.343, %fusion.342), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.70.clone.clone (param_0.1261: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { - %param_0.1261 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %dynamic_slice.318 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1375, %param_2.1167, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.556 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.318), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1089 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.556), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1246 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1703 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1246), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1702 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1089, %mul.1703), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1088 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1088), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.555 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.36.clone.clone (param_0.1247: bf16[4,4096,32,128], param_1.1376: s32[], param_2.1168: f32[4,128], param_3.844: bf16[4,4,128,4096], param_4.525: bf16[4096]) -> bf16[4,128,32,128] { + %param_2.1168 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.844 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) + %param_4.525 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.332 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1168, %param_3.844, %param_1.1376, %param_4.525), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1247 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.331 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1247, %param_1.1376), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.107 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.332, %fusion.331), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.70.clone.clone (param_0.1248: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { + %param_0.1248 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1248), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.129 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1248), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.187 = (bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) } @@ -1502,172 +1502,172 @@ StackFrames %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} %constant.1123 = s32[]{:T(128)} constant(2) %closed_call.83 = s32[64]{0:T(128)} broadcast(%constant.1123), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1699 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1094 = f32[64]{0:T(128)} convert(%mul.1699), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %mul.1693 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1082 = f32[64]{0:T(128)} convert(%mul.1693), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %constant.1122 = f32[]{:T(128)} constant(0.0078125) %closed_call.82 = f32[64]{0:T(128)} broadcast(%constant.1122), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1094, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1082, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.84, %div.843), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} } -%fused_computation.117.clone.clone (param_0.1245: f32[64], param_1.1378: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1378 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1378), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1245 = f32[64]{0:T(128)S(1)} parameter(0) - %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1245), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} +%fused_computation.117.clone.clone (param_0.1232: f32[64], param_1.1365: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1365 = f32[4,128]{1,0:T(4,128)} parameter(1) + %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1365), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %param_0.1232 = f32[64]{0:T(128)S(1)} parameter(0) + %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1232), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %div.844 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.846, %div.845), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1095 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1083 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} %convert_element_type.829.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1095, %convert_element_type.829.clone.3) + ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1083, %convert_element_type.829.clone.3) } -%fused_computation.120.clone.clone (param_0.1252: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1252 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.120.clone.clone (param_0.1239: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1239 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1130 = bf16[]{:T(256)} constant(-inf) - %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1239, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1239, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.61, %pad.60), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.554 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.119.clone.clone (param_0.1246: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1246 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.119.clone.clone (param_0.1233: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1233 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1125 = bf16[]{:T(256)} constant(-inf) - %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1233, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1233, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.44 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.59, %pad.58), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.544 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.73.clone.clone (param_0.1262: bf16[4,128,32,64], param_1.1390: bf16[4,128,32,64], param_2.1178: bf16[4,128,32,128], param_3.849: bf16[4,128,128], param_4.531: bf16[4,128,128]) -> bf16[4,32,128,128] { - %param_2.1178 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.531 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1713 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.531), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1711 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1178, %mul.1713), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1390 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.73.clone.clone (param_0.1249: bf16[4,128,32,64], param_1.1377: bf16[4,128,32,64], param_2.1169: bf16[4,128,32,128], param_3.845: bf16[4,128,128], param_4.526: bf16[4,128,128]) -> bf16[4,32,128,128] { + %param_2.1169 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.526 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1707 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.526), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1705 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1169, %mul.1707), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1377 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1136 = bf16[]{:T(256)} constant(-inf) - %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1390, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1262 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1262, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1377, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1249 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1249, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.47 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.65, %pad.64), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.849 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1712 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.849), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1710 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1712), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1711, %mul.1710), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.562 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.90.clone.clone (param_0.1254: f32[4,128], param_1.1384: bf16[4,4,128,4096], param_2.1173: s32[], param_3.844: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.844 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.844), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1384 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1173 = s32[]{:T(128)S(6)} parameter(2) + %param_3.845 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1706 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.845), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1704 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1706), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1705, %mul.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.557 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.90.clone.clone (param_0.1241: f32[4,128], param_1.1371: bf16[4,4,128,4096], param_2.1164: s32[], param_3.840: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.840 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.422 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.840), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1371 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1164 = s32[]{:T(128)S(6)} parameter(2) %constant.1132 = s32[]{:T(128)} constant(0) - %dynamic_slice.322 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1384, %param_2.1173, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.557 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1099 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1703 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1254), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1702 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1099, %mul.1703), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1098 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1098), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.556 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.425), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.64.clone.1.clone.clone (param_0.1253: bf16[4,4096,8,128], param_1.1383: s32[]) -> bf16[4096,8,128,1] { - %param_0.1253 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1383 = s32[]{:T(128)S(6)} parameter(1) + %dynamic_slice.316 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1371, %param_2.1164, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.316), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1087 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1241 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1697 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1241), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1696 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1087, %mul.1697), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1086 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1696), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.421 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.422, %convert_element_type.1086), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.64.clone.1.clone.clone (param_0.1240: bf16[4,4096,8,128], param_1.1370: s32[]) -> bf16[4096,8,128,1] { + %param_0.1240 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1370 = s32[]{:T(128)S(6)} parameter(1) %constant.1131 = s32[]{:T(128)} constant(0) - %dynamic_slice.321 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1383, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.555 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.315 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1240, %param_1.1370, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.550 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.315), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.89.clone.clone (param_0.1255: bf16[4,4096,8,128], param_1.1385: s32[], param_2.1174: f32[4,128], param_3.845: bf16[4,4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,8,128] { - %param_2.1174 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.845 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) - %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.340 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1174, %param_3.845, %param_1.1385, %param_4.528), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1255 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.341 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1255, %param_1.1385), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.112 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.340, %fusion.341), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.89.clone.clone (param_0.1242: bf16[4,4096,8,128], param_1.1372: s32[], param_2.1165: f32[4,128], param_3.841: bf16[4,4,128,4096], param_4.523: bf16[4096]) -> bf16[4,128,8,128] { + %param_2.1165 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.841 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1372 = s32[]{:T(128)S(6)} parameter(1) + %param_4.523 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.329 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1165, %param_3.841, %param_1.1372, %param_4.523), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1242 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.330 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1242, %param_1.1372), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.329, %fusion.330), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } -%fused_computation.106.clone.clone (param_0.1256: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1256 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.106.clone.clone (param_0.1243: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { + %param_0.1243 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1243), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1243), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.186 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) } -%fused_computation.109.clone.clone (param_0.1257: bf16[4,128,8,64], param_1.1386: bf16[4,128,8,64], param_2.1175: bf16[4,128,8,128], param_3.846: bf16[4,128,128], param_4.529: bf16[4,128,128]) -> bf16[4,8,128,128] { - %param_2.1175 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.529 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1707 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.529), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1705 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1175, %mul.1707), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1386 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.109.clone.clone (param_0.1244: bf16[4,128,8,64], param_1.1373: bf16[4,128,8,64], param_2.1166: bf16[4,128,8,128], param_3.842: bf16[4,128,128], param_4.524: bf16[4,128,128]) -> bf16[4,8,128,128] { + %param_2.1166 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.524 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1701 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.524), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1699 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1166, %mul.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1373 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1133 = bf16[]{:T(256)} constant(-inf) - %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1386, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1257 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1257, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1373, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1244 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1244, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.46 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.63, %pad.62), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.846 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1706 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.846), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1704 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1706), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1705, %mul.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.558 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_3.842 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1700 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1698 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1700), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1699, %mul.1698), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.135.clone.clone (param_0.1248: bf16[4,4096,8,128], param_1.1380: s32[]) -> bf16[1,4096,8,128] { - %param_0.1248 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.135.clone.clone (param_0.1235: bf16[4,4096,8,128], param_1.1367: s32[]) -> bf16[1,4096,8,128] { + %param_0.1235 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) + %param_1.1367 = s32[]{:T(128)S(6)} parameter(1) %constant.1128 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.319 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1248, %param_1.1380, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %dynamic_slice.313 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1235, %param_1.1367, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } -%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1249: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { - %param_0.1249 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.550 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1236: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { + %param_0.1236 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1236), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} + ROOT %bitcast.545 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.88.clone.clone.clone.clone (param_0.1250: f32[4,128], param_1.1381: bf16[4,4,128,4096], param_2.1171: s32[], param_3.842: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.842 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1381 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) +%fused_computation.88.clone.clone.clone.clone (param_0.1237: f32[4,128], param_1.1368: bf16[4,4,128,4096], param_2.1162: s32[], param_3.838: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.838 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.420 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.838), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1368 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1162 = s32[]{:T(128)S(6)} parameter(2) %constant.1129 = s32[]{:T(128)} constant(0) - %dynamic_slice.320 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1381, %param_2.1171, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1097 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1250 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1701 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1250), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1700 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1097, %mul.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1096 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1700), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1096), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.114.clone.clone (param_0.1251: bf16[1,4096,8,128], param_1.1382: f32[4,128], param_2.1172: bf16[4,4,128,4096], param_3.843: s32[], param_4.527: bf16[4096]) -> bf16[4,8,128,128] { - %param_1.1382 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1172 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.843 = s32[]{:T(128)S(6)} parameter(3) - %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.339 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1382, %param_2.1172, %param_3.843, %param_4.527), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1251 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.338 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1251), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.111 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.339, %fusion.338), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.366.clone.clone (param_0.1286: f32[4,32,128,128]) -> (f32[4,32,128,1], f32[4,32,128]) { - %param_0.1286 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1286), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.262.clone.3 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.192 = (f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}, f32[4,32,128]{2,1,0:T(8,128)S(1)}) tuple(%slice.11, %bitcast.262.clone.3) + %dynamic_slice.314 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1368, %param_2.1162, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.547 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.314), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1085 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.547), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1237 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1695 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1237), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1694 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1085, %mul.1695), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1084 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1694), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.419 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.420, %convert_element_type.1084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.546 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.419), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.114.clone.clone (param_0.1238: bf16[1,4096,8,128], param_1.1369: f32[4,128], param_2.1163: bf16[4,4,128,4096], param_3.839: s32[], param_4.522: bf16[4096]) -> bf16[4,8,128,128] { + %param_1.1369 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1163 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.839 = s32[]{:T(128)S(6)} parameter(3) + %param_4.522 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.328 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1369, %param_2.1163, %param_3.839, %param_4.522), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1238 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.327 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1238), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.105 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.328, %fusion.327), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + ROOT %bitcast.548 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.146.clone.clone (param_0.1273: f32[4,32,128,128]) -> (f32[4,32,128], f32[4,32,128,1]) { + %param_0.1273 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) + %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1273), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.570 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} + ROOT %tuple.192 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.570, %slice.11) } %region_13.16 (reduce_sum.120: f32[], reduce_sum.121: f32[]) -> f32[] { @@ -1676,34 +1676,34 @@ StackFrames ROOT %reduce_sum.122 = f32[]{:T(128)} add(%reduce_sum.120, %reduce_sum.121), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1263: bf16[4,32,128,4096], param_1.1391: s32[]) -> bf16[32,128,4096,1] { - %param_0.1263 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1250: bf16[4,32,128,4096], param_1.1378: s32[]) -> bf16[32,128,4096,1] { + %param_0.1250 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1378 = s32[]{:T(128)S(6)} parameter(1) %constant.1137 = s32[]{:T(128)} constant(0) - %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1263, %param_1.1391, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.563 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.319 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1250, %param_1.1378, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.558 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.319), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1264: bf16[4,32,128,128]) -> bf16[4,128,32,128] { - %param_0.1264 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.564 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1251: bf16[4,32,128,128]) -> bf16[4,128,32,128] { + %param_0.1251 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.559 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.61.clone.clone (param_0.1265: bf16[4,32,128,4096], param_1.1392: s32[], param_2.1179: bf16[4,32,128,128], param_3.850: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { - %param_3.850 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1392 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.61.clone.clone (param_0.1252: bf16[4,32,128,4096], param_1.1379: s32[], param_2.1170: bf16[4,32,128,128], param_3.846: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { + %param_3.846 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) %constant.365.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.208.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.850, %param_1.1392, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.208.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1179 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.83.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1179), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1265 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.82.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1265, %param_1.1392), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.62.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.83.clone.3, %fusion.82.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %dynamic_slice.210.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.846, %param_1.1379, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.210.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %param_2.1170 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.80.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1170), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_0.1252 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.79.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1252, %param_1.1379), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.60.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.80.clone.3, %fusion.79.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %add.635.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.207.clone.3, %bitcast.182.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1102 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.215 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1102, %convert_element_type.1102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %convert_element_type.1090 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.215 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1090, %convert_element_type.1090), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} %constant.1138 = f32[]{:T(128)} constant(0) %reduce.177 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.215, %constant.1138), dimensions={2}, to_apply=%region_13.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.188 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.177, %add.635.clone.3) @@ -1715,140 +1715,140 @@ StackFrames ROOT %add.623 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.122.clone.clone (param_0.1247: bf16[4,4096], param_1.1379: s32[]) -> bf16[4096] { - %param_0.1247 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.122.clone.clone (param_0.1234: bf16[4,4096], param_1.1366: s32[]) -> bf16[4096] { + %param_0.1234 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1366 = s32[]{:T(128)S(6)} parameter(1) %constant.1126 = s32[]{:T(128)} constant(0) - %dynamic_slice.318 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1247, %param_1.1379, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %dynamic_slice.312 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1234, %param_1.1366, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} %constant.1127 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.318, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.312, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.12.clone.clone.clone (param_0.1266: bf16[4,14336,4096], param_1.1393: s32[]) -> bf16[14336,4096,1] { - %param_0.1266 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1393 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.12.clone.clone.clone (param_0.1253: bf16[4,14336,4096], param_1.1380: s32[]) -> bf16[14336,4096,1] { + %param_0.1253 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) %constant.1139 = s32[]{:T(128)} constant(0) - %dynamic_slice.326 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1266, %param_1.1393, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.566 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.326), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.320 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1380, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.561 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %bitcast_fusion.3.clone.clone (bitcast_input.12: bf16[4,128,4096]) -> bf16[4,128,4096] { %bitcast_input.12 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.565 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) + ROOT %bitcast.560 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) } -%fused_computation.13.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1394: bf16[4,14336,4096], param_2.1180: s32[]) -> bf16[14336,4,128] { - %param_1.1394 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) - %fusion.344 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1394, %param_2.1180), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %fusion.345 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1267), kind=kLoop, calls=%bitcast_fusion.3.clone.clone - ROOT %convolution.114 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.344, %fusion.345), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.13.clone.clone (param_0.1254: bf16[4,128,4096], param_1.1381: bf16[4,14336,4096], param_2.1171: s32[]) -> bf16[14336,4,128] { + %param_1.1381 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) + %fusion.333 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1381, %param_2.1171), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1254 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.334 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1254), kind=kLoop, calls=%bitcast_fusion.3.clone.clone + ROOT %convolution.108 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.333, %fusion.334), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} } -%fused_computation.144.clone.1.clone (param_0.1268: f32[4,128]) -> f32[4,128] { - %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.144.clone.1.clone (param_0.1255: f32[4,128]) -> f32[4,128] { + %param_0.1255 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1141 = f32[]{:T(128)} constant(0.000244140625) %closed_call.86 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1141), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1255, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1140 = f32[]{:T(128)} constant(1e-05) %closed_call.85 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1140), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.861 = f32[4,128]{1,0:T(4,128)} add(%div.847, %closed_call.85), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.98 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.11.clone.1.clone.clone (param_0.1272: bf16[4,4096,14336], param_1.1398: s32[]) -> bf16[4096,14336,1] { - %param_0.1272 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1398 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.11.clone.1.clone.clone (param_0.1259: bf16[4,4096,14336], param_1.1385: s32[]) -> bf16[4096,14336,1] { + %param_0.1259 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) %constant.1143 = s32[]{:T(128)} constant(0) - %dynamic_slice.328 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1272, %param_1.1398, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.568 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.322 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1259, %param_1.1385, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.563 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.96.clone.2.clone.clone (param_0.1273: f32[4,128], param_1.1399: bf16[4,128,4096], param_2.1183: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1183 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.432 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1183), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1399 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1106 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1717 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1273), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1716 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1106, %mul.1717), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1105 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1716), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.431 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.432, %convert_element_type.1105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.96.clone.2.clone.clone (param_0.1260: f32[4,128], param_1.1386: bf16[4,128,4096], param_2.1174: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1174), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1094 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1260 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1711 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1260), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1710 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1094, %mul.1711), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1093 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1710), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.23.clone.clone (param_0.1274: bf16[4,4096,14336], param_1.1400: s32[], param_2.1184: f32[4,128], param_3.852: bf16[4,128,4096], param_4.533: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1184 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.533 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.349 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1184, %param_3.852, %param_4.533), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1274 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1400 = s32[]{:T(128)S(6)} parameter(1) - %fusion.348 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1274, %param_1.1400), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.116 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.349, %fusion.348), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.23.clone.clone (param_0.1261: bf16[4,4096,14336], param_1.1387: s32[], param_2.1175: f32[4,128], param_3.848: bf16[4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,14336] { + %param_2.1175 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.848 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.338 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1175, %param_3.848, %param_4.528), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1261 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) + %fusion.337 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1261, %param_1.1387), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.110 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.338, %fusion.337), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } -%fused_computation.14.clone.1.clone.clone (param_0.1275: bf16[4,4096,14336], param_1.1401: s32[]) -> bf16[4096,14336,1] { - %param_0.1275 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1401 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.14.clone.1.clone.clone (param_0.1262: bf16[4,4096,14336], param_1.1388: s32[]) -> bf16[4096,14336,1] { + %param_0.1262 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1388 = s32[]{:T(128)S(6)} parameter(1) %constant.1144 = s32[]{:T(128)} constant(0) - %dynamic_slice.329 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1275, %param_1.1401, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.569 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.329), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.323 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1262, %param_1.1388, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.564 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.39.clone.1.clone.clone (param_0.1276: bf16[14336,4,128], param_1.1402: bf16[4,128,14336]) -> bf16[4,128,14336] { - %param_1.1402 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.39.clone.1.clone.clone (param_0.1263: bf16[14336,4,128], param_1.1389: bf16[4,128,14336]) -> bf16[4,128,14336] { + %param_1.1389 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1145 = bf16[]{:T(256)} constant(1) %jit_silu_.44 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1145), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1402), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.69 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.862 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.848 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.862), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1719 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1402, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %param_0.1276 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) - %bitcast.570 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - ROOT %mul.1718 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1719, %bitcast.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1713 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1389, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %param_0.1263 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.565 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + ROOT %mul.1712 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1713, %bitcast.565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} } -%fused_computation.21.clone.clone (param_0.1277: bf16[4,4096,14336], param_1.1403: s32[], param_2.1185: bf16[14336,4,128], param_3.853: bf16[4,128,14336]) -> bf16[4,128,4096] { - %param_2.1185 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %param_3.853 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1185, %param_3.853), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_0.1277 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1403 = s32[]{:T(128)S(6)} parameter(1) - %fusion.350 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1277, %param_1.1403), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.350), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.21.clone.clone (param_0.1264: bf16[4,4096,14336], param_1.1390: s32[], param_2.1176: bf16[14336,4,128], param_3.849: bf16[4,128,14336]) -> bf16[4,128,4096] { + %param_2.1176 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) + %param_3.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1176, %param_3.849), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %param_0.1264 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1390 = s32[]{:T(128)S(6)} parameter(1) + %fusion.339 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1264, %param_1.1390), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.111 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.339), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} } -%fused_computation.14.clone.clone.clone (param_0.1269: bf16[4,4096,14336], param_1.1395: s32[]) -> bf16[4096,14336,1] { - %param_0.1269 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1395 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.14.clone.clone.clone (param_0.1256: bf16[4,4096,14336], param_1.1382: s32[]) -> bf16[4096,14336,1] { + %param_0.1256 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1382 = s32[]{:T(128)S(6)} parameter(1) %constant.1142 = s32[]{:T(128)} constant(0) - %dynamic_slice.327 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1269, %param_1.1395, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.567 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.321 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1256, %param_1.1382, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.562 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.96.clone.1.clone.clone (param_0.1270: f32[4,128], param_1.1396: bf16[4,128,4096], param_2.1181: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1181 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1181), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1396 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1104 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1270 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1715 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1270), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1714 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1104, %mul.1715), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1103 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1714), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.430, %convert_element_type.1103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.96.clone.1.clone.clone (param_0.1257: f32[4,128], param_1.1383: bf16[4,128,4096], param_2.1172: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1172 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1172), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1092 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1257 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1709 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1257), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1708 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1092, %mul.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1091 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1708), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1091), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.20.clone.clone (param_0.1271: bf16[4,4096,14336], param_1.1397: s32[], param_2.1182: f32[4,128], param_3.851: bf16[4,128,4096], param_4.532: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1182 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.532 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.347 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1182, %param_3.851, %param_4.532), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1271 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1397 = s32[]{:T(128)S(6)} parameter(1) - %fusion.346 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1271, %param_1.1397), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.115 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.347, %fusion.346), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.20.clone.clone (param_0.1258: bf16[4,4096,14336], param_1.1384: s32[], param_2.1173: f32[4,128], param_3.847: bf16[4,128,4096], param_4.527: bf16[4096]) -> bf16[4,128,14336] { + %param_2.1173 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.847 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.336 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1173, %param_3.847, %param_4.527), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1258 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1384 = s32[]{:T(128)S(6)} parameter(1) + %fusion.335 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1258, %param_1.1384), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.109 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.336, %fusion.335), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } %region_14.17 (reduce_sum.126: f32[], reduce_sum.127: f32[]) -> f32[] { @@ -1857,63 +1857,63 @@ StackFrames ROOT %reduce_sum.128 = f32[]{:T(128)} add(%reduce_sum.126, %reduce_sum.127), metadata={op_name="checkpoint/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1278: bf16[4,4096,14336], param_1.1404: s32[]) -> bf16[4096,14336,1] { - %param_0.1278 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1404 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1265: bf16[4,4096,14336], param_1.1391: s32[]) -> bf16[4096,14336,1] { + %param_0.1265 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) %constant.1146 = s32[]{:T(128)} constant(0) - %dynamic_slice.330 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1278, %param_1.1404, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.571 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.324 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1265, %param_1.1391, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.566 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1279: bf16[4,128,14336], param_1.1405: bf16[4,128,14336], param_2.1186: bf16[14336,4,128]) -> bf16[4,128,14336] { - %param_2.1186 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %bitcast.572 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1186), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %param_1.1405 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %mul.1724 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.572, %param_1.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} +%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1266: bf16[4,128,14336], param_1.1392: bf16[4,128,14336], param_2.1177: bf16[14336,4,128]) -> bf16[4,128,14336] { + %param_2.1177 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.567 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1177), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + %param_1.1392 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %mul.1718 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.567, %param_1.1392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1147 = bf16[]{:T(256)} constant(1) %jit_silu_.45 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1147), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %param_0.1279 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %param_0.1266 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1266), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.70 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.131), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.863 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.70, %jit_silu_.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.45, %add.863), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1723 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1724, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1722 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1279, %mul.1724), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1717 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1718, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1716 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1266, %mul.1718), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} %sub.98 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} subtract(%jit_silu_.45, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/sub" stack_frame_id=0} - %mul.1721 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1720 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1722, %mul.1721), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1723, %mul.1720), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} -} - -%fused_computation.63.clone.clone (param_0.1280: bf16[4,128,4096], param_1.1406: bf16[4096], param_2.1187: bf16[4,128,4096], param_3.854: bf16[4,4096,14336], param_4.534: s32[], param_5.435: bf16[4,128,14336], param_6.304: bf16[4,128,14336], param_7.200: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { - %param_0.1280 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1108 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_2.1187 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_5.435 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) - %param_6.304 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) - %param_7.200 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) - %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.435, %param_6.304, %param_7.200), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} - %param_3.854 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.534 = s32[]{:T(128)S(6)} parameter(4) - %fusion.79.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.854, %param_4.534), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.60.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.79.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1187, %convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1406 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) - %dot_general.434 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1406), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.433 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.434), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1107 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.433), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1725 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1108, %convert_element_type.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1715 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1714 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1716, %mul.1715), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1717, %mul.1714), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} +} + +%fused_computation.63.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1393: bf16[4096], param_2.1178: bf16[4,128,4096], param_3.850: bf16[4,4096,14336], param_4.529: s32[], param_5.425: bf16[4,128,14336], param_6.291: bf16[4,128,14336], param_7.188: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { + %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1096 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1267), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_2.1178 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_5.425 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) + %param_6.291 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) + %param_7.188 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) + %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.425, %param_6.291, %param_7.188), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} + %param_3.850 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) + %param_4.529 = s32[]{:T(128)S(6)} parameter(4) + %fusion.91.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.850, %param_4.529), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.64.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.91.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1178, %convolution.64.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %param_1.1393 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1393), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.430), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1095 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.429), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} + %mul.1719 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1096, %convert_element_type.1095), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1148 = f32[]{:T(128)} constant(0) - %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1725, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} + %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1719, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.189 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.178, %add_any.132.clone.3) } -%fused_computation.140.clone.clone (param_0.1281: f32[4,128], param_1.1407: f32[4,128]) -> f32[4,128] { - %param_0.1281 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1407 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.140.clone.clone (param_0.1268: f32[4,128], param_1.1394: f32[4,128]) -> f32[4,128] { + %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1394 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.1152 = f32[]{:T(128)} constant(0.000244140625) %closed_call.89 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1152), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1407, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1394, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1151 = f32[]{:T(128)} constant(1e-05) %closed_call.88 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1151), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.864 = f32[4,128]{1,0:T(4,128)} add(%div.851, %closed_call.88), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} @@ -1921,11 +1921,11 @@ StackFrames %div.850 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.99, %add.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1150 = f32[]{:T(128)} constant(-0.5) %closed_call.87 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1150), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1728 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1727 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1281, %mul.1728), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1722 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1721 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %mul.1722), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1149 = f32[]{:T(128)} constant(0.00048828125) - %mul.1729 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - ROOT %mul.1726 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1727, %mul.1729), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1723 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + ROOT %mul.1720 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1721, %mul.1723), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} } %region_20.24 (dot_general.187: bf16[], dot_general.188: bf16[]) -> bf16[] { @@ -1934,29 +1934,29 @@ StackFrames ROOT %add.173 = bf16[]{:T(256)} add(%dot_general.187, %dot_general.188), metadata={op_name="add.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.94.clone.clone (param_0.1282: bf16[4,128,4096], param_1.1408: f32[4,128], param_2.1188: bf16[4,128,4096], param_3.855: bf16[4,128,4096], param_4.535: f32[4,128], param_5.436: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { - %param_0.1282 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_2.1188 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %convert_element_type.1110 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1188), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_1.1408 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1731 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1408), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1730 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1109 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1730), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %convert_element_type.1109), metadata={op_name="multiply.204"} +%fused_computation.94.clone.clone (param_0.1269: bf16[4,128,4096], param_1.1395: f32[4,128], param_2.1179: bf16[4,128,4096], param_3.851: bf16[4,128,4096], param_4.530: f32[4,128], param_5.426: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { + %param_0.1269 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_2.1179 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1098 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1179), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_1.1395 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1725 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1395), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1724 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1098, %mul.1725), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1097 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1724), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1269, %convert_element_type.1097), metadata={op_name="multiply.204"} %constant.1153 = bf16[]{:T(256)} constant(0) %reduce.179 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.271, %constant.1153), dimensions={0,1}, to_apply=%region_20.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.855 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_5.436 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) - %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.436), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_5.426 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.426), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1269, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} %convert_element_type.753.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.263.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1142.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_4.535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %mul.1151.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.535), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %mul.1141.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1151.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1142.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1725), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %param_4.530 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %mul.1151.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.530), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1141.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1098, %mul.1151.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %add_any.126.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1142.clone.3, %mul.1141.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} %convert_element_type.751.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.126.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.855, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.851, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} ROOT %tuple.190 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.179, %add_any.124.clone.3) } @@ -1966,35 +1966,35 @@ StackFrames ROOT %add.169 = f32[]{:T(128)} add(%dot_general.184, %dot_general.185), metadata={op_name="add.31"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1283: bf16[4,32,128,4096], param_1.1409: s32[]) -> bf16[32,128,4096,1] { - %param_0.1283 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1409 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1270: bf16[4,32,128,4096], param_1.1396: s32[]) -> bf16[32,128,4096,1] { + %param_0.1270 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1396 = s32[]{:T(128)S(6)} parameter(1) %constant.1154 = s32[]{:T(128)} constant(0) - %dynamic_slice.331 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1283, %param_1.1409, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.573 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1270, %param_1.1396, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.568 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1284: bf16[4,128,4096]) -> bf16[4,128,4096,1] { - %param_0.1284 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.574 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} +%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1271: bf16[4,128,4096]) -> bf16[4,128,4096,1] { + %param_0.1271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.569 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1271), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} } -%fused_computation.66.clone.clone (param_0.1285: bf16[4,32,128,128], param_1.1410: bf16[4,32,128,4096], param_2.1189: s32[], param_3.856: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { - %param_0.1285 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} - %param_3.856 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %fusion.95.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.856), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1410 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1189 = s32[]{:T(128)S(6)} parameter(2) - %fusion.94.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1410, %param_2.1189), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.64.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.95.clone.3, %fusion.94.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.66.clone.clone (param_0.1272: bf16[4,32,128,128], param_1.1397: bf16[4,32,128,4096], param_2.1180: s32[], param_3.852: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { + %param_0.1272 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1272), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} + %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %fusion.84.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.852), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %param_1.1397 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) + %fusion.83.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1397, %param_2.1180), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.62.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.84.clone.3, %fusion.83.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} %constant.619.clone.3 = bf16[]{:T(256)} constant(0.25) %div.442.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.619.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} - %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.64.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} + %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.62.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} %bitcast.209.clone.3 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%div.441.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} %convert.123 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.209.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert.1" stack_frame_id=0} %multiply.272 = f32[4,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.124, %convert.123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/multiply" stack_frame_id=0} %constant.1155 = f32[]{:T(128)} constant(0) - %dot_general.435 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} - ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.435, %bitcast.209.clone.3) + %dot_general.431 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} + ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.431, %bitcast.209.clone.3) } diff --git a/tests/utils/reference_hlo_qwen3_1.7b.txt b/tests/utils/reference_hlo_qwen3_1.7b.txt index f1ede66966..b54b810621 100644 --- a/tests/utils/reference_hlo_qwen3_1.7b.txt +++ b/tests/utils/reference_hlo_qwen3_1.7b.txt @@ -32,7 +32,7 @@ StackFrames %param_1.5 = s32[512]{0:T(512)S(1)} parameter(1) %reshape.451 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.5), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} %transpose.466 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.451), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)} parameter(2) %reshape.452 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} reshape(%param_2.4), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} %transpose.467 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} transpose(%reshape.452), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} ROOT %scatter.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.466, %transpose.467), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_42.47.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} @@ -50,43 +50,43 @@ StackFrames ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.386, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.277 (param_0.1368: f32[151936,2048], param_1.1556: f32[], param_2.1314: f32[], param_3.918: f32[], param_4.556: f32[151936,2048], param_5.468: f32[], param_6.358: bf16[151936,2048], param_7.201: bf16[151936,2048,1], param_8.118: pred[], param_9.97: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { - %param_0.1368 = f32[151936,2048]{1,0:T(8,128)} parameter(0) +%fused_computation.277 (param_0.1367: f32[151936,2048], param_1.1549: f32[], param_2.1311: f32[], param_3.918: f32[], param_4.554: f32[151936,2048], param_5.467: f32[], param_6.356: bf16[151936,2048], param_7.196: bf16[151936,2048,1], param_8.113: pred[], param_9.94: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { + %param_0.1367 = f32[151936,2048]{1,0:T(8,128)} parameter(0) %param_3.918 = f32[]{:T(128)S(6)} parameter(3) %mul.1926.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_3.918), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.118 = pred[]{:T(512)S(6)} parameter(8) - %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.118), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_7.201 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) - %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.201), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %param_8.113 = pred[]{:T(512)S(6)} parameter(8) + %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.113), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_7.196 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) + %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.196), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %convert_element_type.1409.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.464.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_6.358 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_6.356 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.356), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %add_any.197.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1409.clone.1, %convert_element_type.1408.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} - %param_5.468 = f32[]{:T(128)} parameter(5) - %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_5.467 = f32[]{:T(128)} parameter(5) + %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.467), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.859.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add_any.197.clone.1, %div.860.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.267.clone.1 = f32[151936,2048]{1,0:T(8,128)} select(%select_n.268.clone.1, %add_any.197.clone.1, %div.859.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1092.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.844.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1092.clone.1), dimensions={}, metadata={op_name="broadcast.74"} %mul.1932.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_9.97 = f32[151936,2048]{1,0:T(8,128)} parameter(9) + %param_9.94 = f32[151936,2048]{1,0:T(8,128)} parameter(9) %constant.1096.clone.1 = f32[]{:T(128)} constant(0.9) %mul.1933.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1096.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1931.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.97, %mul.1933.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1931.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.94, %mul.1933.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.941.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1932.clone.1, %mul.1931.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) - %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1311 = f32[]{:T(128)S(6)} parameter(2) + %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1311), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.65.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %select_n.267.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1095.clone.1 = f32[]{:T(128)} constant(0.05) %mul.1930.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1095.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.1928.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %mul.1930.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.556 = f32[151936,2048]{1,0:T(8,128)} parameter(4) + %param_4.554 = f32[151936,2048]{1,0:T(8,128)} parameter(4) %constant.1094.clone.1 = f32[]{:T(128)} constant(0.95) %mul.1929.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1094.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1927.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.556, %mul.1929.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1927.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.554, %mul.1929.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.940.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1928.clone.1, %mul.1927.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) - %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1549 = f32[]{:T(128)S(6)} parameter(1) + %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1549), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.854.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.940.clone.1, %div.855.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.62.clone.1 = f32[151936,2048]{1,0:T(8,128)} sqrt(%div.854.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1093.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -94,14 +94,14 @@ StackFrames %add.938.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%sqrt.62.clone.1, %add.939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.426.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%div.856.clone.1, %add.938.clone.1), metadata={op_name="multiply.61"} %div.853.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.941.clone.1, %multiply.426.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1925.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1368, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1925.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1367, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.937.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%div.853.clone.1, %mul.1925.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1924.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%mul.1926.clone.1, %add.937.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1368, %mul.1924.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1367, %mul.1924.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.214 = f32[151936,2048]{1,0:T(8,128)} multiply(%add.936.clone.1, %add.936.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1200 = f32[]{:T(128)} constant(0) - %reduce.176 = f32[]{:T(128)} reduce(%square.214, %constant.1200), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1198 = f32[]{:T(128)} constant(0) + %reduce.176 = f32[]{:T(128)} reduce(%square.214, %constant.1198), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1198), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.144 = (f32[]{:T(128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.176, %add.936.clone.1, %add.940.clone.1, %add.941.clone.1, %reduce.178.clone.1) } @@ -111,64 +111,64 @@ StackFrames ROOT %reduce_sum.319 = f32[]{:T(128)} add(%reduce_sum.317, %reduce_sum.318), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.367.clone.clone (param_0.1355: f32[4,128], param_1.1549: bf16[4,128,2048], param_2.1290: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1290 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.480 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1290), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1549 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1451 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1355 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2083 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1355), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2082 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1451, %mul.2083), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.480, %convert_element_type.1450), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.367.clone.clone (param_0.1354: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1287: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1287 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1287), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1445 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1354 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2075 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1354), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.2074 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1445, %mul.2075), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1444 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2074), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.478 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.479, %convert_element_type.1444), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.289.clone.clone.clone (param_0.1356: bf16[4,128,151936], param_1.1550: s32[4,128], param_2.1291: f32[4,128], param_3.911: f32[4,128], param_4.546: bf16[4,128], param_5.446: f32[4,128]) -> bf16[4,128,151936] { - %param_5.446 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2087 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.446), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.289.clone.clone.clone (param_0.1355: bf16[4,128,151936], param_1.1543: s32[4,128], param_2.1288: f32[4,128], param_3.911: f32[4,128], param_4.544: bf16[4,128], param_5.445: f32[4,128]) -> bf16[4,128,151936] { + %param_5.445 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.2079 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.445), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2086 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1356 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1454 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1356), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.546 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1454, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %mul.2078 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1355 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1448 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1355), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.544 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.544), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1448, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.62 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2085 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2086, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1291 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1291), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2085, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1550 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.2077 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2078, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2077, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1543), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.48 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.47 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1453 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1453), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2084 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2087, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1452 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.281 (param_0.1381: bf16[151936,2048], param_1.1569: f32[4,128], param_2.1327: bf16[4,128,2048], param_3.931: bf16[2048], param_4.569: bf16[4,128,151936], param_5.481: s32[4,128], param_6.371: f32[4,128], param_7.214: f32[4,128], param_8.131: bf16[4,128], param_9.98: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { - %param_4.569 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) - %param_5.481 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.371 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.214 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) - %param_8.131 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) - %param_9.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.569, %param_5.481, %param_6.371, %param_7.214, %param_8.131, /*index=5*/%param_9.98), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1327 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1447 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1447), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2076 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2079, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1446 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2076), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.281 (param_0.1380: bf16[151936,2048], param_1.1562: f32[4,128], param_2.1324: bf16[4,128,2048], param_3.931: bf16[2048], param_4.567: bf16[4,128,151936], param_5.480: s32[4,128], param_6.369: f32[4,128], param_7.209: f32[4,128], param_8.126: bf16[4,128], param_9.95: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { + %param_4.567 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) + %param_5.480 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.369 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.209 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) + %param_8.126 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) + %param_9.95 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.567, %param_5.480, %param_6.369, %param_7.209, %param_8.126, /*index=5*/%param_9.95), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1562 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1324 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) %param_3.931 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1569, %param_2.1327, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.269.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %fusion.268.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1562, %param_2.1324, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.268.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %bitcast.333 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%convolution.86.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %convert_element_type.1323 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_0.1381 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_0.1380 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1380), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %add_any.184 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1323, %convert_element_type.1322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} %square.215 = f32[151936,2048]{1,0:T(8,128)} multiply(%add_any.184, %add_any.184), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1213 = f32[]{:T(128)} constant(0) - %reduce.177 = f32[]{:T(128)} reduce(%square.215, %constant.1213), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1211 = f32[]{:T(128)} constant(0) + %reduce.177 = f32[]{:T(128)} reduce(%square.215, %constant.1211), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.166 = (f32[]{:T(128)}, bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.177, %convolution.86.clone.1) } @@ -178,23 +178,23 @@ StackFrames ROOT %reduce_sum.394 = f32[]{:T(128)} add(%reduce_sum.389, %reduce_sum.393), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.288 (param_0.1392: bf16[4,128,151936], param_1.1577: f32[4,128], param_2.1330: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { - %param_2.1330 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1330), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.288 (param_0.1391: bf16[4,128,151936], param_1.1570: f32[4,128], param_2.1327: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { + %param_2.1327 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1327), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.25 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.24 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_0.1391 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} %param_3.933 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) %sub.73 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.933), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.64 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1340, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1577 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1577), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1570 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1570), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.60 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.1225 = f32[]{:T(128)} constant(0) - %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1225), dimensions={}, metadata={op_name="broadcast.109"} + %constant.1223 = f32[]{:T(128)} constant(0) + %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.109"} %mul.1765 = f32[4,128,151936]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.769), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1765, %constant.1225), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1765, %constant.1223), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_9.12 (reduce_sum.186: f32[], reduce_sum.190: f32[]) -> f32[] { @@ -203,15 +203,15 @@ StackFrames ROOT %reduce_sum.191 = f32[]{:T(128)} add(%reduce_sum.186, %reduce_sum.190), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.293 (param_0.1393: bf16[4,128,151936], param_1.1578: bf16[4,128]) -> f32[4,128] { - %param_0.1393 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1578 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1578), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.293 (param_0.1392: bf16[4,128,151936], param_1.1571: bf16[4,128]) -> f32[4,128] { + %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1571 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1571), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.70 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1346, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.54 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.1226 = f32[]{:T(128)} constant(0) - ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1226), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %constant.1224 = f32[]{:T(128)} constant(0) + ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1224), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_33.38 (reduce_sum.269: f32[], reduce_sum.270: f32[]) -> f32[] { @@ -220,12 +220,12 @@ StackFrames ROOT %reduce_sum.274 = f32[]{:T(128)} add(%reduce_sum.269, %reduce_sum.270), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.298 (param_0.1387: f32[4,6144,2048]) -> f32[] { - %param_0.1387 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) - %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.298 (param_0.1386: f32[4,6144,2048]) -> f32[] { + %param_0.1386 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) + %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.218 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%bitcast.347, %bitcast.347), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1219 = f32[]{:T(128)} constant(0) - ROOT %reduce.181 = f32[]{:T(128)} reduce(%square.218, %constant.1219), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1217 = f32[]{:T(128)} constant(0) + ROOT %reduce.181 = f32[]{:T(128)} reduce(%square.218, %constant.1217), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_32.37 (reduce_sum.263: f32[], reduce_sum.267: f32[]) -> f32[] { @@ -240,35 +240,35 @@ StackFrames ROOT %reduce_sum.262 = f32[]{:T(128)} add(%reduce_sum.260, %reduce_sum.261), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.300 (param_0.1388: f32[4,2048,6144], param_1.1573: f32[4,2048,6144]) -> (f32[], f32[]) { - %param_0.1388 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) - %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.300 (param_0.1387: f32[4,2048,6144], param_1.1566: f32[4,2048,6144]) -> (f32[], f32[]) { + %param_0.1387 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) + %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.221 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.351, %bitcast.351), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1220 = f32[]{:T(128)} constant(0) - %reduce.182 = f32[]{:T(128)} reduce(%square.221, %constant.1220), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1573 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) - %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %constant.1218 = f32[]{:T(128)} constant(0) + %reduce.182 = f32[]{:T(128)} reduce(%square.221, %constant.1218), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1566 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) + %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1566), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.224.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.355.clone.1, %bitcast.355.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.183.clone.1 = f32[]{:T(128)} reduce(%square.224.clone.1, %constant.1220), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.183.clone.1 = f32[]{:T(128)} reduce(%square.224.clone.1, %constant.1218), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.167 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.182, %reduce.183.clone.1) } -%fused_computation.303 (param_0.901: f32[6144,4,2048]) -> bf16[4,6144,2048] { - %param_0.901 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) - %copy.190 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.901), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.190), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.303 (param_0.900: f32[6144,4,2048]) -> bf16[4,6144,2048] { + %param_0.900 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) + %copy.192 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.900), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} + ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.304 (param_0.903: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.903 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.191 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.903), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.304 (param_0.902: f32[2048,4,6144]) -> bf16[4,2048,6144] { + %param_0.902 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.193 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.902), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} + ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.305 (param_0.905: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.905 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.192 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.905), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.305 (param_0.904: f32[2048,4,6144]) -> bf16[4,2048,6144] { + %param_0.904 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.194 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.904), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} + ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_62.67 (reduce_sum.416: f32[], reduce_sum.417: f32[]) -> f32[] { @@ -283,39 +283,39 @@ StackFrames ROOT %reduce_sum.340 = f32[]{:T(128)} add(%reduce_sum.338, %reduce_sum.339), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.306 (param_0.1377: f32[6144,4,2048], param_1.1565: f32[], param_2.1323: f32[], param_3.927: f32[], param_4.565: f32[6144,4,2048], param_5.477: f32[], param_6.367: f32[4,6144,2048], param_7.210: pred[], param_8.127: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { - %param_0.1377 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) +%fused_computation.306 (param_0.1376: f32[6144,4,2048], param_1.1558: f32[], param_2.1320: f32[], param_3.927: f32[], param_4.563: f32[6144,4,2048], param_5.476: f32[], param_6.365: f32[4,6144,2048], param_7.205: pred[], param_8.122: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { + %param_0.1376 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) %param_3.927 = f32[]{:T(128)S(6)} parameter(3) %mul.1998.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_3.927), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.210 = pred[]{:T(512)S(6)} parameter(7) - %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.210), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.367 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) - %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.477 = f32[]{:T(128)} parameter(5) - %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.205 = pred[]{:T(512)S(6)} parameter(7) + %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.365 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) + %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.476 = f32[]{:T(128)} parameter(5) + %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.931.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%bitcast.482.clone.1, %div.932.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.303.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} select(%select_n.304.clone.1, %bitcast.482.clone.1, %div.931.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1146.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.886.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1146.clone.1), dimensions={}, metadata={op_name="broadcast.83"} %mul.2004.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.127 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) + %param_8.122 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) %constant.1150.clone.1 = f32[]{:T(128)} constant(0.9) %mul.2005.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1150.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2003.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.127, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2003.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.122, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.989.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2004.clone.1, %mul.2003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) - %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) + %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.74.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %select_n.303.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1149.clone.1 = f32[]{:T(128)} constant(0.05) %mul.2002.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1149.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.2000.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%integer_pow.74.clone.1, %mul.2002.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.565 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) + %param_4.563 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) %constant.1148.clone.1 = f32[]{:T(128)} constant(0.95) %mul.2001.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1148.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1999.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.565, %mul.2001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1999.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.563, %mul.2001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.988.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2000.clone.1, %mul.1999.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1565 = f32[]{:T(128)S(6)} parameter(1) - %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1565), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) + %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.926.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.988.clone.1, %div.927.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.71.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} sqrt(%div.926.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1147.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -323,14 +323,14 @@ StackFrames %add.986.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%sqrt.71.clone.1, %add.987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.435.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%div.928.clone.1, %add.986.clone.1), metadata={op_name="multiply.52"} %div.925.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.989.clone.1, %multiply.435.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1997.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1997.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1376, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.985.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%div.925.clone.1, %mul.1997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1996.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%mul.1998.clone.1, %add.985.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1377, %mul.1996.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1376, %mul.1996.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.225 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%add.984.clone.1, %add.984.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1209 = f32[]{:T(128)} constant(0) - %reduce.184 = f32[]{:T(128)} reduce(%square.225, %constant.1209), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1207 = f32[]{:T(128)} constant(0) + %reduce.184 = f32[]{:T(128)} reduce(%square.225, %constant.1207), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1207), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.145 = (f32[]{:T(128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.184, %add.984.clone.1, %add.988.clone.1, %add.989.clone.1, %reduce.187.clone.1) } @@ -346,39 +346,39 @@ StackFrames ROOT %reduce_sum.337 = f32[]{:T(128)} add(%reduce_sum.332, %reduce_sum.333), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.307 (param_0.1378: f32[2048,4,6144], param_1.1566: f32[], param_2.1324: f32[], param_3.928: f32[], param_4.566: f32[2048,4,6144], param_5.478: f32[], param_6.368: f32[4,2048,6144], param_7.211: pred[], param_8.128: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) +%fused_computation.307 (param_0.1377: f32[2048,4,6144], param_1.1559: f32[], param_2.1321: f32[], param_3.928: f32[], param_4.564: f32[2048,4,6144], param_5.477: f32[], param_6.366: f32[4,2048,6144], param_7.206: pred[], param_8.123: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { + %param_0.1377 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) %param_3.928 = f32[]{:T(128)S(6)} parameter(3) %mul.2008.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.928), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.211 = pred[]{:T(512)S(6)} parameter(7) - %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.211), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.368 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.478 = f32[]{:T(128)} parameter(5) - %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.206 = pred[]{:T(512)S(6)} parameter(7) + %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.366 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.477 = f32[]{:T(128)} parameter(5) + %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.939.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.484.clone.1, %div.940.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.307.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.308.clone.1, %bitcast.484.clone.1, %div.939.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1152.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.892.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1152.clone.1), dimensions={}, metadata={op_name="broadcast.85"} %mul.2012.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.128 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) + %param_8.123 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) %constant.1156.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.891.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1156.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2011.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.128, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2011.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.123, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.994.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2012.clone.1, %mul.2011.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1324 = f32[]{:T(128)S(6)} parameter(2) - %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1324), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) + %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.75.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %select_n.307.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1155.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.890.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1155.clone.1), dimensions={}, metadata={op_name="broadcast.73"} %mul.2010.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.75.clone.1, %broadcast.890.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.566 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) + %param_4.564 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) %constant.1154.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.889.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1154.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2009.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.566, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2009.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.564, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.993.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2010.clone.1, %mul.2009.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1566 = f32[]{:T(128)S(6)} parameter(1) - %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1566), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) + %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.934.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.993.clone.1, %div.935.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.72.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.934.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1153.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -386,14 +386,14 @@ StackFrames %add.992.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.72.clone.1, %broadcast.887.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.436.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.936.clone.1, %add.992.clone.1), metadata={op_name="multiply.51"} %div.933.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.994.clone.1, %multiply.436.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2007.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2007.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.991.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.933.clone.1, %mul.2007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.2006.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2008.clone.1, %add.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2006.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1377, %mul.2006.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.226 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.990.clone.1, %add.990.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1210 = f32[]{:T(128)} constant(0) - %reduce.185 = f32[]{:T(128)} reduce(%square.226, %constant.1210), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1210), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1208 = f32[]{:T(128)} constant(0) + %reduce.185 = f32[]{:T(128)} reduce(%square.226, %constant.1208), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1208), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.146 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.185, %add.990.clone.1, %add.993.clone.1, %add.994.clone.1, %reduce.188.clone.1) } @@ -409,39 +409,39 @@ StackFrames ROOT %reduce_sum.331 = f32[]{:T(128)} add(%reduce_sum.326, %reduce_sum.330), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.308 (param_0.1379: f32[2048,4,6144], param_1.1567: f32[], param_2.1325: f32[], param_3.929: f32[], param_4.567: f32[2048,4,6144], param_5.479: f32[], param_6.369: f32[4,2048,6144], param_7.212: pred[], param_8.129: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1379 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) +%fused_computation.308 (param_0.1378: f32[2048,4,6144], param_1.1560: f32[], param_2.1322: f32[], param_3.929: f32[], param_4.565: f32[2048,4,6144], param_5.478: f32[], param_6.367: f32[4,2048,6144], param_7.207: pred[], param_8.124: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { + %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) %param_3.929 = f32[]{:T(128)S(6)} parameter(3) %mul.2015.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.929), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.212 = pred[]{:T(512)S(6)} parameter(7) - %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.212), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.369 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.369), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.479 = f32[]{:T(128)} parameter(5) - %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.207 = pred[]{:T(512)S(6)} parameter(7) + %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.367 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.478 = f32[]{:T(128)} parameter(5) + %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.947.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.486.clone.1, %div.948.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.311.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.312.clone.1, %bitcast.486.clone.1, %div.947.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1158.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.898.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1158.clone.1), dimensions={}, metadata={op_name="broadcast.85"} %mul.2019.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.129 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) + %param_8.124 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) %constant.1162.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.897.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1162.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2018.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.129, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2018.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.124, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.999.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2019.clone.1, %mul.2018.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1325 = f32[]{:T(128)S(6)} parameter(2) - %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1325), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) + %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.76.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %select_n.311.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1161.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.896.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1161.clone.1), dimensions={}, metadata={op_name="broadcast.73"} %mul.2017.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.76.clone.1, %broadcast.896.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.567 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) + %param_4.565 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) %constant.1160.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.895.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1160.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2016.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.567, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2016.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.565, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.998.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2017.clone.1, %mul.2016.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1567 = f32[]{:T(128)S(6)} parameter(1) - %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1567), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) + %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.942.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.998.clone.1, %div.943.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.73.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.942.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1159.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -449,14 +449,14 @@ StackFrames %add.997.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.73.clone.1, %broadcast.893.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.437.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.944.clone.1, %add.997.clone.1), metadata={op_name="multiply.50"} %div.941.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.999.clone.1, %multiply.437.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2014.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1379, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2014.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.996.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.941.clone.1, %mul.2014.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.2013.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2015.clone.1, %add.996.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1379, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.227 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.995.clone.1, %add.995.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1211 = f32[]{:T(128)} constant(0) - %reduce.186 = f32[]{:T(128)} reduce(%square.227, %constant.1211), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1211), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1209 = f32[]{:T(128)} constant(0) + %reduce.186 = f32[]{:T(128)} reduce(%square.227, %constant.1209), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.147 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.186, %add.995.clone.1, %add.998.clone.1, %add.999.clone.1, %reduce.189.clone.1) } @@ -466,12 +466,12 @@ StackFrames ROOT %reduce_sum.304 = f32[]{:T(128)} add(%reduce_sum.302, %reduce_sum.303), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.324 (param_0.1382: f32[4,2048,16,128]) -> f32[] { - %param_0.1382 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.324 (param_0.1381: f32[4,2048,16,128]) -> f32[] { + %param_0.1381 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.230 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%bitcast.362, %bitcast.362), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1214 = f32[]{:T(128)} constant(0) - ROOT %reduce.190 = f32[]{:T(128)} reduce(%square.230, %constant.1214), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1212 = f32[]{:T(128)} constant(0) + ROOT %reduce.190 = f32[]{:T(128)} reduce(%square.230, %constant.1212), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_38.43 (reduce_sum.296: f32[], reduce_sum.297: f32[]) -> f32[] { @@ -480,18 +480,18 @@ StackFrames ROOT %reduce_sum.298 = f32[]{:T(128)} add(%reduce_sum.296, %reduce_sum.297), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.326 (param_0.1383: f32[4,16,128,2048]) -> f32[] { - %param_0.1383 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.326 (param_0.1382: f32[4,16,128,2048]) -> f32[] { + %param_0.1382 = f32[4,16,128,2048]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.233 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%bitcast.366, %bitcast.366), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1215 = f32[]{:T(128)} constant(0) - ROOT %reduce.191 = f32[]{:T(128)} reduce(%square.233, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1213 = f32[]{:T(128)} constant(0) + ROOT %reduce.191 = f32[]{:T(128)} reduce(%square.233, %constant.1213), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.327 (param_0.950: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { - %param_0.950 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) - %copy.193 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.950), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.327 (param_0.949: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { + %param_0.949 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) + %copy.195 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.949), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} + ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.195), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_68.73 (reduce_sum.449: f32[], reduce_sum.450: f32[]) -> f32[] { @@ -506,39 +506,39 @@ StackFrames ROOT %reduce_sum.373 = f32[]{:T(128)} add(%reduce_sum.368, %reduce_sum.372), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.328 (param_0.1371: f32[2048,4,16,128], param_1.1559: f32[], param_2.1317: f32[], param_3.921: f32[], param_4.559: f32[2048,4,16,128], param_5.471: f32[], param_6.361: f32[4,2048,16,128], param_7.204: pred[], param_8.121: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { - %param_0.1371 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.328 (param_0.1370: f32[2048,4,16,128], param_1.1552: f32[], param_2.1314: f32[], param_3.921: f32[], param_4.557: f32[2048,4,16,128], param_5.470: f32[], param_6.359: f32[4,2048,16,128], param_7.199: pred[], param_8.116: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { + %param_0.1370 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.921 = f32[]{:T(128)S(6)} parameter(3) %mul.1950.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_3.921), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.204 = pred[]{:T(512)S(6)} parameter(7) - %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.361 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.471 = f32[]{:T(128)} parameter(5) - %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.199 = pred[]{:T(512)S(6)} parameter(7) + %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.199), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.359 = f32[4,2048,16,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.470 = f32[]{:T(128)} parameter(5) + %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.883.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%bitcast.470.clone.1, %div.884.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.279.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} select(%select_n.280.clone.1, %bitcast.470.clone.1, %div.883.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1110.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.858.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1110.clone.1), dimensions={}, metadata={op_name="broadcast.75"} %mul.1956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.121 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) + %param_8.116 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1114.clone.1 = f32[]{:T(128)} constant(0.9) %mul.1957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1114.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.121, %mul.1957.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.116, %mul.1957.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1956.clone.1, %mul.1955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) - %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) + %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.68.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %select_n.279.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1113.clone.1 = f32[]{:T(128)} constant(0.05) %mul.1954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1113.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.1952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.68.clone.1, %mul.1954.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.559 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) + %param_4.557 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1112.clone.1 = f32[]{:T(128)} constant(0.95) %mul.1953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1112.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1951.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.559, %mul.1953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1951.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %mul.1953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1952.clone.1, %mul.1951.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) - %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1552 = f32[]{:T(128)S(6)} parameter(1) + %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1552), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.878.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.956.clone.1, %div.879.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.65.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} sqrt(%div.878.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1111.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -546,14 +546,14 @@ StackFrames %add.954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%sqrt.65.clone.1, %add.955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.429.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%div.880.clone.1, %add.954.clone.1), metadata={op_name="multiply.58"} %div.877.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.957.clone.1, %multiply.429.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1949.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1949.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1370, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%div.877.clone.1, %mul.1949.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1948.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%mul.1950.clone.1, %add.953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.1948.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1370, %mul.1948.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.234 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%add.952.clone.1, %add.952.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1203 = f32[]{:T(128)} constant(0) - %reduce.192 = f32[]{:T(128)} reduce(%square.234, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1201 = f32[]{:T(128)} constant(0) + %reduce.192 = f32[]{:T(128)} reduce(%square.234, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.148 = (f32[]{:T(128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.192, %add.952.clone.1, %add.956.clone.1, %add.957.clone.1, %reduce.194.clone.1) } @@ -569,39 +569,39 @@ StackFrames ROOT %reduce_sum.367 = f32[]{:T(128)} add(%reduce_sum.365, %reduce_sum.366), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.329 (param_0.1372: f32[16,4,128,2048], param_1.1560: f32[], param_2.1318: f32[], param_3.922: f32[], param_4.560: f32[16,4,128,2048], param_5.472: f32[], param_6.362: f32[4,16,128,2048], param_7.205: pred[], param_8.122: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { - %param_0.1372 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.329 (param_0.1371: f32[16,4,128,2048], param_1.1553: f32[], param_2.1315: f32[], param_3.922: f32[], param_4.558: f32[16,4,128,2048], param_5.471: f32[], param_6.360: f32[4,16,128,2048], param_7.200: pred[], param_8.117: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { + %param_0.1371 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) %param_3.922 = f32[]{:T(128)S(6)} parameter(3) %mul.1960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_3.922), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.205 = pred[]{:T(512)S(6)} parameter(7) - %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.362 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.472 = f32[]{:T(128)} parameter(5) - %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.200 = pred[]{:T(512)S(6)} parameter(7) + %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.200), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.360 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.471 = f32[]{:T(128)} parameter(5) + %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.891.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%bitcast.472.clone.1, %div.892.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.283.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} select(%select_n.284.clone.1, %bitcast.472.clone.1, %div.891.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1116.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.860.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1116.clone.1), dimensions={}, metadata={op_name="broadcast.76"} %mul.1966.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.122 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) + %param_8.117 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) %constant.1120.clone.1 = f32[]{:T(128)} constant(0.9) %mul.1967.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1120.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1965.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.122, %mul.1967.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1965.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.117, %mul.1967.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1966.clone.1, %mul.1965.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) - %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) + %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.69.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %select_n.283.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1119.clone.1 = f32[]{:T(128)} constant(0.05) %mul.1964.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1119.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.1962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%integer_pow.69.clone.1, %mul.1964.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.560 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) + %param_4.558 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) %constant.1118.clone.1 = f32[]{:T(128)} constant(0.95) %mul.1963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1118.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.560, %mul.1963.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.558, %mul.1963.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1962.clone.1, %mul.1961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) - %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1553 = f32[]{:T(128)S(6)} parameter(1) + %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1553), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.886.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.962.clone.1, %div.887.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.66.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} sqrt(%div.886.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1117.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -609,14 +609,14 @@ StackFrames %add.960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%sqrt.66.clone.1, %add.961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.430.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%div.888.clone.1, %add.960.clone.1), metadata={op_name="multiply.57"} %div.885.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.963.clone.1, %multiply.430.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1372, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%div.885.clone.1, %mul.1959.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%mul.1960.clone.1, %add.959.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1372, %mul.1958.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.1958.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.235 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%add.958.clone.1, %add.958.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1204 = f32[]{:T(128)} constant(0) - %reduce.193 = f32[]{:T(128)} reduce(%square.235, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1202 = f32[]{:T(128)} constant(0) + %reduce.193 = f32[]{:T(128)} reduce(%square.235, %constant.1202), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1202), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.149 = (f32[]{:T(128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.193, %add.958.clone.1, %add.962.clone.1, %add.963.clone.1, %reduce.195.clone.1) } @@ -632,23 +632,23 @@ StackFrames ROOT %reduce_sum.289 = f32[]{:T(128)} add(%reduce_sum.284, %reduce_sum.288), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.341 (param_0.1385: f32[4,2048,8,128], param_1.1571: f32[4,2048,8,128]) -> (f32[], f32[]) { - %param_0.1385 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) - %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.341 (param_0.1384: f32[4,2048,8,128], param_1.1564: f32[4,2048,8,128]) -> (f32[], f32[]) { + %param_0.1384 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.238 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1217 = f32[]{:T(128)} constant(0) - %reduce.196 = f32[]{:T(128)} reduce(%square.238, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1571 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(1) - %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %constant.1215 = f32[]{:T(128)} constant(0) + %reduce.196 = f32[]{:T(128)} reduce(%square.238, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1564 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) + %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.241.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.197.clone.1 = f32[]{:T(128)} reduce(%square.241.clone.1, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.197.clone.1 = f32[]{:T(128)} reduce(%square.241.clone.1, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.168 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.196, %reduce.197.clone.1) } -%fused_computation.344 (param_0.982: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.982 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %copy.194 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.982), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.344 (param_0.981: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { + %param_0.981 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.196 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.981), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} + ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.196), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_70.75 (reduce_sum.458: f32[], reduce_sum.459: f32[]) -> f32[] { @@ -663,39 +663,39 @@ StackFrames ROOT %reduce_sum.382 = f32[]{:T(128)} add(%reduce_sum.380, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.345 (param_0.1369: f32[2048,4,8,128], param_1.1557: f32[], param_2.1315: f32[], param_3.919: f32[], param_4.557: f32[2048,4,8,128], param_5.469: f32[], param_6.359: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1369 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.345 (param_0.1368: f32[2048,4,8,128], param_1.1550: f32[], param_2.1312: f32[], param_3.919: f32[], param_4.555: f32[2048,4,8,128], param_5.468: f32[], param_6.357: f32[4,2048,8,128], param_7.197: pred[], param_8.114: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { + %param_0.1368 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.919 = f32[]{:T(128)S(6)} parameter(3) %mul.1936.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.919), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.202 = pred[]{:T(512)S(6)} parameter(7) - %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.359 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.469 = f32[]{:T(128)} parameter(5) - %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.197 = pred[]{:T(512)S(6)} parameter(7) + %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.357 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.357), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.468 = f32[]{:T(128)} parameter(5) + %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.867.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.466.clone.1, %div.868.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.271.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.272.clone.1, %bitcast.466.clone.1, %div.867.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1098.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.850.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1098.clone.1), dimensions={}, metadata={op_name="broadcast.80"} %mul.1940.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %param_8.114 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1102.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.849.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1102.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.1939.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1939.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.114, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.946.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1940.clone.1, %mul.1939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) - %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1312 = f32[]{:T(128)S(6)} parameter(2) + %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1312), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.66.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %select_n.271.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1101.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.848.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1101.clone.1), dimensions={}, metadata={op_name="broadcast.69"} %mul.1938.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.66.clone.1, %broadcast.848.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.557 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %param_4.555 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1100.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.847.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1100.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1937.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1937.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.555, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.945.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1938.clone.1, %mul.1937.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) - %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1550 = f32[]{:T(128)S(6)} parameter(1) + %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.862.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.945.clone.1, %div.863.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.63.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.862.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1099.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -703,14 +703,14 @@ StackFrames %add.944.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.845.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.427.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.864.clone.1, %add.944.clone.1), metadata={op_name="multiply.60"} %div.861.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.946.clone.1, %multiply.427.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1935.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1369, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1935.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1368, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.943.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.861.clone.1, %mul.1935.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1934.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1936.clone.1, %add.943.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1369, %mul.1934.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1368, %mul.1934.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.242 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.942.clone.1, %add.942.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1201 = f32[]{:T(128)} constant(0) - %reduce.198 = f32[]{:T(128)} reduce(%square.242, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1199 = f32[]{:T(128)} constant(0) + %reduce.198 = f32[]{:T(128)} reduce(%square.242, %constant.1199), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1199), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.150 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.198, %add.942.clone.1, %add.945.clone.1, %add.946.clone.1, %reduce.200.clone.1) } @@ -726,39 +726,39 @@ StackFrames ROOT %reduce_sum.358 = f32[]{:T(128)} add(%reduce_sum.353, %reduce_sum.354), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.346 (param_0.1374: f32[2048,4,8,128], param_1.1562: f32[], param_2.1320: f32[], param_3.924: f32[], param_4.562: f32[2048,4,8,128], param_5.474: f32[], param_6.364: f32[4,2048,8,128], param_7.207: pred[], param_8.124: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1374 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.346 (param_0.1373: f32[2048,4,8,128], param_1.1555: f32[], param_2.1317: f32[], param_3.924: f32[], param_4.560: f32[2048,4,8,128], param_5.473: f32[], param_6.362: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { + %param_0.1373 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) %param_3.924 = f32[]{:T(128)S(6)} parameter(3) %mul.1977.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.924), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.207 = pred[]{:T(512)S(6)} parameter(7) - %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.364 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.474 = f32[]{:T(128)} parameter(5) - %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.202 = pred[]{:T(512)S(6)} parameter(7) + %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.362 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.473 = f32[]{:T(128)} parameter(5) + %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.907.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.476.clone.1, %div.908.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.291.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.292.clone.1, %bitcast.476.clone.1, %div.907.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1128.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.872.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1128.clone.1), dimensions={}, metadata={op_name="broadcast.80"} %mul.1981.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.124 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1132.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.871.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1132.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.1980.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.124, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1980.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.973.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1981.clone.1, %mul.1980.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) - %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) + %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.71.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %select_n.291.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1131.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.870.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1131.clone.1), dimensions={}, metadata={op_name="broadcast.69"} %mul.1979.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.71.clone.1, %broadcast.870.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.562 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %param_4.560 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1130.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.869.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1130.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1978.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.562, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1978.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.560, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.972.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1979.clone.1, %mul.1978.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1562 = f32[]{:T(128)S(6)} parameter(1) - %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1562), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1555 = f32[]{:T(128)S(6)} parameter(1) + %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1555), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.902.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.972.clone.1, %div.903.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.68.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.902.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1129.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -766,30 +766,30 @@ StackFrames %add.971.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.68.clone.1, %broadcast.867.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.432.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.904.clone.1, %add.971.clone.1), metadata={op_name="multiply.55"} %div.901.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.973.clone.1, %multiply.432.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1976.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1374, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1976.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1373, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.970.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.901.clone.1, %mul.1976.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1975.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1977.clone.1, %add.970.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1374, %mul.1975.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1373, %mul.1975.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.243 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.969.clone.1, %add.969.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1206 = f32[]{:T(128)} constant(0) - %reduce.199 = f32[]{:T(128)} reduce(%square.243, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) + %constant.1204 = f32[]{:T(128)} constant(0) + %reduce.199 = f32[]{:T(128)} reduce(%square.243, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) } -%fused_computation.362 (param_0.1056: bf16[4,128,2048], param_1.1117: f32[4,128], param_2.830: f32[4,128], param_3.495: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { - %param_3.495 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) +%fused_computation.362 (param_0.1055: bf16[4,128,2048], param_1.1114: f32[4,128], param_2.829: f32[4,128], param_3.497: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { + %param_3.497 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) %param_4.296 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.448 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.438 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.495, %dot_general.448), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.438), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.830 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1851 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.830), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %dot_general.451 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.441 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.497, %dot_general.451), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.441), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_2.829 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %mul.1851 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.829), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %mul.1843 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1363, %mul.1851), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.1056 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1056), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1117 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1850 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1117), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_0.1055 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1055), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1850 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1114), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %mul.1849 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1374, %mul.1850), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %add_any.193 = f32[4,128,2048]{2,1,0:T(8,128)} add(%mul.1843, %mul.1849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} ROOT %convert_element_type.1361 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.193), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} @@ -801,12 +801,12 @@ StackFrames ROOT %reduce_sum.185 = f32[]{:T(128)} add(%reduce_sum.171, %reduce_sum.184), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.363 (param_0.1394: bf16[4,128,2048]) -> f32[4,128] { - %param_0.1394 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1394), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.363 (param_0.1393: bf16[4,128,2048]) -> f32[4,128] { + %param_0.1393 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %square.246 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1365, %convert_element_type.1365), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} - %constant.1227 = f32[]{:T(128)} constant(0) - ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.246, %constant.1227), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + %constant.1225 = f32[]{:T(128)} constant(0) + ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.246, %constant.1225), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_12.15 (reduce_sum.198: f32[], reduce_sum.199: f32[]) -> f32[] { @@ -815,17 +815,17 @@ StackFrames ROOT %reduce_sum.200 = f32[]{:T(128)} add(%reduce_sum.198, %reduce_sum.199), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.365 (param_0.1389: bf16[4,128,2048], param_1.1574: bf16[4,128,2048], param_2.1328: bf16[2048]) -> f32[4,128] { - %param_0.1389 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1389), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1328 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.447 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1328), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1574, %dot_general.447), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.437), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.365 (param_0.1388: bf16[4,128,2048], param_1.1567: bf16[4,128,2048], param_2.1325: bf16[2048]) -> f32[4,128] { + %param_0.1388 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1388), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1325 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1325), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.440 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1567, %dot_general.450), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.440), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %mul.1847 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1372, %convert_element_type.1371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %constant.1221 = f32[]{:T(128)} constant(0) - ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1847, %constant.1221), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + %constant.1219 = f32[]{:T(128)} constant(0) + ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1847, %constant.1219), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_10.13 (dot_general.190: bf16[], dot_general.191: bf16[]) -> bf16[] { @@ -834,51 +834,51 @@ StackFrames ROOT %add.419 = bf16[]{:T(256)} add(%dot_general.190, %dot_general.191), metadata={op_name="add.82"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.285.clone.clone (param_0.1351: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1351 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%fused_computation.285.clone.clone (param_0.1350: bf16[151936,2048]) -> bf16[151936,2048,1] { + %param_0.1350 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.530 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1350), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} } -%fused_computation.289.clone.1.clone.clone (param_0.1352: bf16[4,128,151936], param_1.1546: s32[4,128], param_2.1285: f32[4,128], param_3.906: f32[4,128], param_4.542: bf16[4,128], param_5.442: f32[4,128]) -> bf16[4,128,151936] { - %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2075 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.442), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.289.clone.1.clone.clone (param_0.1351: bf16[4,128,151936], param_1.1539: s32[4,128], param_2.1282: f32[4,128], param_3.906: f32[4,128], param_4.540: bf16[4,128], param_5.441: f32[4,128]) -> bf16[4,128,151936] { + %param_5.441 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.2067 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.441), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.906 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2074 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1352 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1444 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1352), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.542 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.542), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1444, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %mul.2066 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1351 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1438 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.540 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.540), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1438, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.60 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2073 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2074, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1285 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1285), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2073, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1546 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.2065 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2066, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1282 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1282), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2065, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1539 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.42 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.41 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1443 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1443), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2072 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2075, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1442 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2072), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convert_element_type.1437 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1437), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2064 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2067, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1436 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2064), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} } -%fused_computation.366 (param_0.1350: f32[4,128], param_1.1545: bf16[4,128,2048], param_2.1286: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.543: s32[4,128], param_5.443: f32[4,128], param_6.340: f32[4,128], param_7.199: bf16[4,128], param_8.116: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { +%fused_computation.366 (param_0.1349: f32[4,128], param_1.1538: bf16[4,128,2048], param_2.1283: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.541: s32[4,128], param_5.442: f32[4,128], param_6.338: f32[4,128], param_7.194: bf16[4,128], param_8.111: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { %param_3.907 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.443 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.340 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.199 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.116 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.543, %param_5.443, %param_6.340, %param_7.199, /*index=5*/%param_8.116), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1286 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) - %fusion.251.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1286), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.251.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %param_1.1545 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1545), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1350 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1862 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1350), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_4.541 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.338 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.194 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) + %param_8.111 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) + %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.541, %param_5.442, %param_6.338, %param_7.194, /*index=5*/%param_8.111), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_2.1283 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) + %fusion.250.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1283), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.250.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %param_1.1538 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1349 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1862 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1349), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %mul.1861 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1384, %mul.1862), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %convert_element_type.1383 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1861), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %multiply.420 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%convolution.84.clone.1, %convert_element_type.1383), metadata={op_name="multiply.362"} @@ -887,11 +887,11 @@ StackFrames ROOT %tuple.165 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.204, %convolution.84.clone.1) } -%fused_computation.374 (param_0.1088: f32[64], param_1.1150: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1150 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1150), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.1088 = f32[64]{0:T(128)S(1)} parameter(0) - %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1088), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} +%fused_computation.374 (param_0.1087: f32[64], param_1.1147: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1147 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1147), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %param_0.1087 = f32[64]{0:T(128)S(1)} parameter(0) + %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1087), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %div.717 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.720, %div.718), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.717), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} %convert_element_type.1392 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} @@ -900,19 +900,19 @@ StackFrames ROOT %tuple.158 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1392, %convert_element_type.1391.clone.1) } -%fused_computation.375 (param_0.1085: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1085 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.375 (param_0.1084: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.1084 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1042 = bf16[]{:T(256)} constant(-inf) - %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1084, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1084, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.42 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.46, %pad.45), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } -%fused_computation.376 (param_0.1087: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1087 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.376 (param_0.1086: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.1086 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1041 = bf16[]{:T(256)} constant(-inf) - %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1086, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1086, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.43 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.48, %pad.47), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } @@ -928,16 +928,16 @@ StackFrames ROOT %reduce_sum.277 = f32[]{:T(128)} add(%reduce_sum.275, %reduce_sum.276), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.380 (param_0.1386: f32[4,2048], param_1.1572: f32[4,2048]) -> (f32[], f32[]) { - %param_0.1386 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.380 (param_0.1385: f32[4,2048], param_1.1565: f32[4,2048]) -> (f32[], f32[]) { + %param_0.1385 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.249 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.404, %bitcast.404), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1218 = f32[]{:T(128)} constant(0) - %reduce.205 = f32[]{:T(128)} reduce(%square.249, %constant.1218), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1572 = f32[4,2048]{1,0:T(4,128)} parameter(1) - %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %constant.1216 = f32[]{:T(128)} constant(0) + %reduce.205 = f32[]{:T(128)} reduce(%square.249, %constant.1216), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1565 = f32[4,2048]{1,0:T(4,128)} parameter(1) + %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.252.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.408.clone.1, %bitcast.408.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.206.clone.1 = f32[]{:T(128)} reduce(%square.252.clone.1, %constant.1218), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.206.clone.1 = f32[]{:T(128)} reduce(%square.252.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.169 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.205, %reduce.206.clone.1) } @@ -953,39 +953,39 @@ StackFrames ROOT %reduce_sum.352 = f32[]{:T(128)} add(%reduce_sum.347, %reduce_sum.351), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.383 (param_0.1375: f32[2048,4], param_1.1563: f32[], param_2.1321: f32[], param_3.925: f32[], param_4.563: f32[2048,4], param_5.475: f32[], param_6.365: f32[4,2048], param_7.208: pred[], param_8.125: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.383 (param_0.1374: f32[2048,4], param_1.1556: f32[], param_2.1318: f32[], param_3.925: f32[], param_4.561: f32[2048,4], param_5.474: f32[], param_6.363: f32[4,2048], param_7.203: pred[], param_8.120: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1374 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.925 = f32[]{:T(128)S(6)} parameter(3) %mul.1984.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.925), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.208 = pred[]{:T(512)S(6)} parameter(7) - %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.365 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.475 = f32[]{:T(128)} parameter(5) - %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.203 = pred[]{:T(512)S(6)} parameter(7) + %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.363 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.474 = f32[]{:T(128)} parameter(5) + %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.915.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.478.clone.1, %div.916.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.295.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.296.clone.1, %bitcast.478.clone.1, %div.915.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1134.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.878.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1134.clone.1), dimensions={}, metadata={op_name="broadcast.82"} %mul.1988.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.125 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %param_8.120 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1138.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.877.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1138.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.1987.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.125, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1987.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.978.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1988.clone.1, %mul.1987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) - %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) + %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.72.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %select_n.295.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1137.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.876.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1137.clone.1), dimensions={}, metadata={op_name="broadcast.71"} %mul.1986.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.72.clone.1, %broadcast.876.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.563 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %param_4.561 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1136.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.875.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1136.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1985.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.563, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1985.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.977.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1986.clone.1, %mul.1985.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1563 = f32[]{:T(128)S(6)} parameter(1) - %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1563), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) + %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.910.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.977.clone.1, %div.911.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.69.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.910.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1135.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -993,14 +993,14 @@ StackFrames %add.976.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.69.clone.1, %broadcast.873.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.433.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.912.clone.1, %add.976.clone.1), metadata={op_name="multiply.54"} %div.909.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.978.clone.1, %multiply.433.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1983.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1983.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1374, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.975.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.909.clone.1, %mul.1983.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1982.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1984.clone.1, %add.975.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.1982.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1374, %mul.1982.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.253 = f32[2048,4]{0,1:T(4,128)} multiply(%add.974.clone.1, %add.974.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1207 = f32[]{:T(128)} constant(0) - %reduce.207 = f32[]{:T(128)} reduce(%square.253, %constant.1207), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1207), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1205 = f32[]{:T(128)} constant(0) + %reduce.207 = f32[]{:T(128)} reduce(%square.253, %constant.1205), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.152 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.207, %add.974.clone.1, %add.977.clone.1, %add.978.clone.1, %reduce.209.clone.1) } @@ -1016,39 +1016,39 @@ StackFrames ROOT %reduce_sum.346 = f32[]{:T(128)} add(%reduce_sum.344, %reduce_sum.345), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.384 (param_0.1376: f32[2048,4], param_1.1564: f32[], param_2.1322: f32[], param_3.926: f32[], param_4.564: f32[2048,4], param_5.476: f32[], param_6.366: f32[4,2048], param_7.209: pred[], param_8.126: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1376 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.384 (param_0.1375: f32[2048,4], param_1.1557: f32[], param_2.1319: f32[], param_3.926: f32[], param_4.562: f32[2048,4], param_5.475: f32[], param_6.364: f32[4,2048], param_7.204: pred[], param_8.121: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.926 = f32[]{:T(128)S(6)} parameter(3) %mul.1991.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.926), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.209 = pred[]{:T(512)S(6)} parameter(7) - %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.209), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.366 = f32[4,2048]{1,0:T(4,128)} parameter(6) - %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.476 = f32[]{:T(128)} parameter(5) - %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.204 = pred[]{:T(512)S(6)} parameter(7) + %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.364 = f32[4,2048]{1,0:T(4,128)} parameter(6) + %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.475 = f32[]{:T(128)} parameter(5) + %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.923.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.480.clone.1, %div.924.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.299.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.300.clone.1, %bitcast.480.clone.1, %div.923.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1140.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.884.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1140.clone.1), dimensions={}, metadata={op_name="broadcast.82"} %mul.1995.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.126 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %param_8.121 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1144.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.883.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1144.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.1994.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.126, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1994.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.121, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.983.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1995.clone.1, %mul.1994.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) - %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) + %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.73.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %select_n.299.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1143.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.882.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1143.clone.1), dimensions={}, metadata={op_name="broadcast.71"} %mul.1993.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.73.clone.1, %broadcast.882.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.564 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %param_4.562 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1142.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.881.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1142.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1992.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.564, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1992.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.562, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.982.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1993.clone.1, %mul.1992.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1564 = f32[]{:T(128)S(6)} parameter(1) - %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1564), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) + %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.918.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.982.clone.1, %div.919.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.70.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.918.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1141.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1056,14 +1056,14 @@ StackFrames %add.981.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.70.clone.1, %broadcast.879.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.434.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.920.clone.1, %add.981.clone.1), metadata={op_name="multiply.53"} %div.917.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.983.clone.1, %multiply.434.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1990.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1376, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1990.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.980.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.917.clone.1, %mul.1990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1989.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1991.clone.1, %add.980.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1376, %mul.1989.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.1989.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.254 = f32[2048,4]{0,1:T(4,128)} multiply(%add.979.clone.1, %add.979.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1208 = f32[]{:T(128)} constant(0) - %reduce.208 = f32[]{:T(128)} reduce(%square.254, %constant.1208), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1208), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1206 = f32[]{:T(128)} constant(0) + %reduce.208 = f32[]{:T(128)} reduce(%square.254, %constant.1206), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1206), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.153 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.208, %add.979.clone.1, %add.982.clone.1, %add.983.clone.1, %reduce.210.clone.1) } @@ -1073,12 +1073,12 @@ StackFrames ROOT %reduce_sum.197 = f32[]{:T(128)} add(%reduce_sum.192, %reduce_sum.193), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.395 (param_0.1390: bf16[2048]) -> f32[] { - %param_0.1390 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.395 (param_0.1389: bf16[2048]) -> f32[] { + %param_0.1389 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %square.257 = f32[2048]{0:T(1024)} multiply(%convert_element_type.1396, %convert_element_type.1396), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1222 = f32[]{:T(128)} constant(0) - ROOT %reduce.211 = f32[]{:T(128)} reduce(%square.257, %constant.1222), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1220 = f32[]{:T(128)} constant(0) + ROOT %reduce.211 = f32[]{:T(128)} reduce(%square.257, %constant.1220), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_59.64 (reduce_sum.401: f32[], reduce_sum.402: f32[]) -> f32[] { @@ -1093,39 +1093,39 @@ StackFrames ROOT %reduce_sum.325 = f32[]{:T(128)} add(%reduce_sum.323, %reduce_sum.324), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.396 (param_0.1380: f32[2048], param_1.1568: f32[], param_2.1326: f32[], param_3.930: f32[], param_4.568: f32[2048], param_5.480: f32[], param_6.370: bf16[2048], param_7.213: pred[], param_8.130: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { - %param_0.1380 = f32[2048]{0:T(1024)S(1)} parameter(0) +%fused_computation.396 (param_0.1379: f32[2048], param_1.1561: f32[], param_2.1323: f32[], param_3.930: f32[], param_4.566: f32[2048], param_5.479: f32[], param_6.368: bf16[2048], param_7.208: pred[], param_8.125: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { + %param_0.1379 = f32[2048]{0:T(1024)S(1)} parameter(0) %param_3.930 = f32[]{:T(128)S(6)} parameter(3) %mul.2022.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_3.930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.213 = pred[]{:T(512)S(6)} parameter(7) - %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.213), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.370 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.480 = f32[]{:T(128)} parameter(5) - %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.480), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.208 = pred[]{:T(512)S(6)} parameter(7) + %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.368 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_5.479 = f32[]{:T(128)} parameter(5) + %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.955.clone.1 = f32[2048]{0:T(1024)} divide(%convert_element_type.1411.clone.1, %div.956.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.315.clone.1 = f32[2048]{0:T(1024)} select(%select_n.316.clone.1, %convert_element_type.1411.clone.1, %div.955.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1164.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.900.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1164.clone.1), dimensions={}, metadata={op_name="broadcast.86"} %mul.2028.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.130 = f32[2048]{0:T(1024)S(1)} parameter(8) + %param_8.125 = f32[2048]{0:T(1024)S(1)} parameter(8) %constant.1168.clone.1 = f32[]{:T(128)} constant(0.9) %mul.2029.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1168.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2027.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.130, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2027.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.125, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.1005.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2028.clone.1, %mul.2027.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1326 = f32[]{:T(128)S(6)} parameter(2) - %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1326), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) + %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.77.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %select_n.315.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1167.clone.1 = f32[]{:T(128)} constant(0.05) %mul.2026.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1167.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.2024.clone.1 = f32[2048]{0:T(1024)} multiply(%integer_pow.77.clone.1, %mul.2026.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.568 = f32[2048]{0:T(1024)S(1)} parameter(4) + %param_4.566 = f32[2048]{0:T(1024)S(1)} parameter(4) %constant.1166.clone.1 = f32[]{:T(128)} constant(0.95) %mul.2025.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1166.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2023.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.568, %mul.2025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2023.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.566, %mul.2025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.1004.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2024.clone.1, %mul.2023.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) - %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) + %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.950.clone.1 = f32[2048]{0:T(1024)} divide(%add.1004.clone.1, %div.951.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.74.clone.1 = f32[2048]{0:T(1024)} sqrt(%div.950.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1165.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1133,37 +1133,37 @@ StackFrames %add.1002.clone.1 = f32[2048]{0:T(1024)} add(%sqrt.74.clone.1, %add.1003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.438.clone.1 = f32[2048]{0:T(1024)} multiply(%div.952.clone.1, %add.1002.clone.1), metadata={op_name="multiply.49"} %div.949.clone.1 = f32[2048]{0:T(1024)} divide(%add.1005.clone.1, %multiply.438.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2021.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1380, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2021.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1379, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.1001.clone.1 = f32[2048]{0:T(1024)} add(%div.949.clone.1, %mul.2021.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.2020.clone.1 = f32[2048]{0:T(1024)} multiply(%mul.2022.clone.1, %add.1001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1380, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1379, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.258 = f32[2048]{0:T(1024)} multiply(%add.1000.clone.1, %add.1000.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1212 = f32[]{:T(128)} constant(0) - %reduce.212 = f32[]{:T(128)} reduce(%square.258, %constant.1212), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1212), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1210 = f32[]{:T(128)} constant(0) + %reduce.212 = f32[]{:T(128)} reduce(%square.258, %constant.1210), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1210), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.156 = (f32[]{:T(128)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.212, %add.1000.clone.1, %add.1004.clone.1, %add.1005.clone.1, %reduce.213.clone.1) } -%fused_computation.402 (param_0.1150: s32[512]) -> s32[1024] { +%fused_computation.402 (param_0.1149: s32[512]) -> s32[1024] { %constant.972 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.815 = s32[1024]{0:T(1024)} broadcast(%constant.972), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.1150 = s32[512]{0:T(512)S(1)} parameter(0) + %param_0.1149 = s32[512]{0:T(512)S(1)} parameter(0) %constant.973 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1150, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1149, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %constant.971 = s32[] constant(151935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.814 = s32[1024]{0:T(1024)} broadcast(%constant.971), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.815, %pad.49, %broadcast.814), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } -%fused_computation.405 (param_0.1149: s32[4,128]) -> s32[512] { - %param_0.1149 = s32[4,128]{1,0:T(4,128)} parameter(0) +%fused_computation.405 (param_0.1148: s32[4,128]) -> s32[512] { + %param_0.1148 = s32[4,128]{1,0:T(4,128)} parameter(0) %constant.1065 = s32[]{:T(128)} constant(0) %broadcast.834 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1065), dimensions={}, metadata={op_name="broadcast.95"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1149, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} + %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1148, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} %constant.1051 = s32[]{:T(128)} constant(151936) %add.925 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1051), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1149, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1149), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} + %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1148, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} + %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1148), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} ROOT %bitcast.409 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} } @@ -1179,16 +1179,16 @@ StackFrames ROOT %reduce_sum.295 = f32[]{:T(128)} add(%reduce_sum.290, %reduce_sum.291), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.407 (param_0.1384: f32[4,128], param_1.1570: f32[4,128]) -> (f32[], f32[]) { - %param_0.1384 = f32[4,128]{1,0:T(4,128)} parameter(0) - %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.407 (param_0.1383: f32[4,128], param_1.1563: f32[4,128]) -> (f32[], f32[]) { + %param_0.1383 = f32[4,128]{1,0:T(4,128)} parameter(0) + %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.261 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.413, %bitcast.413), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1216 = f32[]{:T(128)} constant(0) - %reduce.214 = f32[]{:T(128)} reduce(%square.261, %constant.1216), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1570 = f32[4,128]{1,0:T(4,128)} parameter(1) - %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %constant.1214 = f32[]{:T(128)} constant(0) + %reduce.214 = f32[]{:T(128)} reduce(%square.261, %constant.1214), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1563 = f32[4,128]{1,0:T(4,128)} parameter(1) + %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.264.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.417.clone.1, %bitcast.417.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.215.clone.1 = f32[]{:T(128)} reduce(%square.264.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.215.clone.1 = f32[]{:T(128)} reduce(%square.264.clone.1, %constant.1214), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.170 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.214, %reduce.215.clone.1) } @@ -1204,27 +1204,27 @@ StackFrames ROOT %reduce_sum.400 = f32[]{:T(128)} add(%reduce_sum.395, %reduce_sum.396), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.410 (param_0.1391: bf16[4,128], param_1.1576: f32[4,128], param_2.1329: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { +%fused_computation.410 (param_0.1390: bf16[4,128], param_1.1569: f32[4,128], param_2.1326: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { %param_3.932 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) %constant.1170.clone.1 = s32[]{:T(128)} constant(0) %broadcast.901.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1170.clone.1), dimensions={}, metadata={op_name="broadcast.95"} %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.932, %broadcast.901.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1576 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1576), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1391 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1569), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_0.1390 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1390), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} %add.927 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} %square.269 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %add.927), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} - %constant.1224 = f32[]{:T(128)} constant(0) - %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1224), dimensions={}, metadata={op_name="broadcast.99"} + %constant.1222 = f32[]{:T(128)} constant(0) + %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1222), dimensions={}, metadata={op_name="broadcast.99"} %mul.1913 = f32[4,128]{1,0:T(4,128)} multiply(%square.269, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} %mul.1893 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1913, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.216 = f32[]{:T(128)} reduce(%mul.1893, %constant.1224), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1329 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1329), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %reduce.216 = f32[]{:T(128)} reduce(%mul.1893, %constant.1222), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_2.1326 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1326), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} %add.904.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1913), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} %mul.1894.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.904.clone.1, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1894.clone.1, %constant.1224), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1894.clone.1, %constant.1222), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} %mul.1911.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %broadcast.831), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %constant.1068.clone.1 = f32[]{:T(128)} constant(1) %add.922.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1068.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} @@ -1244,39 +1244,39 @@ StackFrames ROOT %reduce_sum.379 = f32[]{:T(128)} add(%reduce_sum.374, %reduce_sum.375), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.411 (param_0.1370: f32[128,4], param_1.1558: f32[], param_2.1316: f32[], param_3.920: f32[], param_4.558: f32[128,4], param_5.470: f32[], param_6.360: f32[4,128], param_7.203: pred[], param_8.120: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1370 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.411 (param_0.1369: f32[128,4], param_1.1551: f32[], param_2.1313: f32[], param_3.920: f32[], param_4.556: f32[128,4], param_5.469: f32[], param_6.358: f32[4,128], param_7.198: pred[], param_8.115: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1369 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.920 = f32[]{:T(128)S(6)} parameter(3) %mul.1943.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.920), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.203 = pred[]{:T(512)S(6)} parameter(7) - %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.360 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.470 = f32[]{:T(128)} parameter(5) - %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.198 = pred[]{:T(512)S(6)} parameter(7) + %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.198), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.358 = f32[4,128]{1,0:T(4,128)} parameter(6) + %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.469 = f32[]{:T(128)} parameter(5) + %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.875.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.468.clone.1, %div.876.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.275.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.276.clone.1, %bitcast.468.clone.1, %div.875.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1104.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.856.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1104.clone.1), dimensions={}, metadata={op_name="broadcast.78"} %mul.1947.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.120 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %param_8.115 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1108.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.855.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1108.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.1946.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1946.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.115, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.951.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1947.clone.1, %mul.1946.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) - %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1313 = f32[]{:T(128)S(6)} parameter(2) + %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1313), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.67.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %select_n.275.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1107.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.854.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1107.clone.1), dimensions={}, metadata={op_name="broadcast.67"} %mul.1945.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.854.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.558 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %param_4.556 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1106.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.853.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1106.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1944.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.558, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1944.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.556, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.950.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1945.clone.1, %mul.1944.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) - %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1551 = f32[]{:T(128)S(6)} parameter(1) + %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1551), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.870.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.950.clone.1, %div.871.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.64.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.870.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1105.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1284,14 +1284,14 @@ StackFrames %add.949.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.851.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.428.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.872.clone.1, %add.949.clone.1), metadata={op_name="multiply.59"} %div.869.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.951.clone.1, %multiply.428.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1942.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1370, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1942.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1369, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.948.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.869.clone.1, %mul.1942.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1941.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1943.clone.1, %add.948.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1370, %mul.1941.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1369, %mul.1941.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.265 = f32[128,4]{0,1:T(4,128)} multiply(%add.947.clone.1, %add.947.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1202 = f32[]{:T(128)} constant(0) - %reduce.217 = f32[]{:T(128)} reduce(%square.265, %constant.1202), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1202), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1200 = f32[]{:T(128)} constant(0) + %reduce.217 = f32[]{:T(128)} reduce(%square.265, %constant.1200), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.159 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.217, %add.947.clone.1, %add.950.clone.1, %add.951.clone.1, %reduce.221.clone.1) } @@ -1307,39 +1307,39 @@ StackFrames ROOT %reduce_sum.361 = f32[]{:T(128)} add(%reduce_sum.359, %reduce_sum.360), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.412 (param_0.1373: f32[128,4], param_1.1561: f32[], param_2.1319: f32[], param_3.923: f32[], param_4.561: f32[128,4], param_5.473: f32[], param_6.363: f32[4,128], param_7.206: pred[], param_8.123: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1373 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.412 (param_0.1372: f32[128,4], param_1.1554: f32[], param_2.1316: f32[], param_3.923: f32[], param_4.559: f32[128,4], param_5.472: f32[], param_6.361: f32[4,128], param_7.201: pred[], param_8.118: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1372 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.923 = f32[]{:T(128)S(6)} parameter(3) %mul.1970.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.923), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.206 = pred[]{:T(512)S(6)} parameter(7) - %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.363 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.473 = f32[]{:T(128)} parameter(5) - %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.201 = pred[]{:T(512)S(6)} parameter(7) + %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.201), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.361 = f32[4,128]{1,0:T(4,128)} parameter(6) + %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.472 = f32[]{:T(128)} parameter(5) + %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.899.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.474.clone.1, %div.900.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.287.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.288.clone.1, %bitcast.474.clone.1, %div.899.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1122.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.866.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1122.clone.1), dimensions={}, metadata={op_name="broadcast.78"} %mul.1974.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.123 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %param_8.118 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1126.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.865.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1126.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.1973.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.123, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1973.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.118, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.968.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1974.clone.1, %mul.1973.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) - %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) + %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.70.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %select_n.287.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1125.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.864.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1125.clone.1), dimensions={}, metadata={op_name="broadcast.67"} %mul.1972.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.864.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.561 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %param_4.559 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1124.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.863.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1124.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1971.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1971.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.559, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.967.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1972.clone.1, %mul.1971.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) - %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1554 = f32[]{:T(128)S(6)} parameter(1) + %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1554), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.894.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.967.clone.1, %div.895.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.67.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.894.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1123.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1347,23 +1347,23 @@ StackFrames %add.966.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.67.clone.1, %broadcast.861.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.431.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.896.clone.1, %add.966.clone.1), metadata={op_name="multiply.56"} %div.893.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.968.clone.1, %multiply.431.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1969.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1373, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1969.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1372, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.965.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.893.clone.1, %mul.1969.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1968.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1970.clone.1, %add.965.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1373, %mul.1968.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1372, %mul.1968.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.266 = f32[128,4]{0,1:T(4,128)} multiply(%add.964.clone.1, %add.964.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1205 = f32[]{:T(128)} constant(0) - %reduce.218 = f32[]{:T(128)} reduce(%square.266, %constant.1205), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1203 = f32[]{:T(128)} constant(0) + %reduce.218 = f32[]{:T(128)} reduce(%square.266, %constant.1203), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1203), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.160 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.218, %add.964.clone.1, %add.967.clone.1, %add.968.clone.1, %reduce.222.clone.1) } -%fused_computation.421 (param_0.1201: f32[4,128], param_1.1323: f32[4,128]) -> f32[4,128] { - %param_0.1201 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1323 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.421 (param_0.1200: f32[4,128], param_1.1320: f32[4,128]) -> f32[4,128] { + %param_0.1200 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1320 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.1045 = f32[]{:T(128)} constant(0.00048828125) %broadcast.837 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1045), dimensions={}, metadata={op_name="broadcast.399"} - %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1323, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1320, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1043 = f32[]{:T(128)} constant(1e-06) %add.935 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1043), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.934 = f32[4,128]{1,0:T(4,128)} add(%div.767, %add.935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} @@ -1372,7 +1372,7 @@ StackFrames %constant.1040 = f32[]{:T(128)} constant(-0.5) %mul.1919 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1040), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %mul.1910 = f32[4,128]{1,0:T(4,128)} multiply(%div.754, %mul.1919), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1909 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1201, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1909 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1200, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.1039 = f32[]{:T(128)} constant(0.0009765625) %mul.1918 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1039), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} ROOT %mul.1908 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1909, %mul.1918), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} @@ -1384,31 +1384,31 @@ StackFrames ROOT %reduce_sum.139 = s32[]{:T(128)} add(%reduce_sum.137, %reduce_sum.138), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.425 (param_0.1220: pred[4,128]) -> s32[] { - %param_0.1220 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1220), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} +%fused_computation.425 (param_0.1219: pred[4,128]) -> s32[] { + %param_0.1219 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1219), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} %constant.1066 = s32[]{:T(128)} constant(0) ROOT %reduce.220 = s32[]{:T(128)} reduce(%convert_element_type.1403, %constant.1066), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.428 (param_0.1203: f32[4,128]) -> f32[4,128] { - %param_0.1203 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.428 (param_0.1202: f32[4,128]) -> f32[4,128] { + %param_0.1202 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1046 = f32[]{:T(128)} constant(0.00048828125) %broadcast.829 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1046), dimensions={}, metadata={op_name="broadcast.399"} - %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1203, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1202, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1044 = f32[]{:T(128)} constant(1e-06) %add.924 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1044), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.921 = f32[4,128]{1,0:T(4,128)} add(%div.759, %add.924), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} ROOT %rsqrt.166 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.921), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} } -%fused_computation.429 (param_0.1204: pred[4,128], param_1.1575: f32[]) -> f32[4,128] { - %param_0.1204 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1575 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1575), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} - %constant.1223 = f32[]{:T(128)} constant(0) - %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.99"} - ROOT %mul.1920 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1204, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.429 (param_0.1203: pred[4,128], param_1.1568: f32[]) -> f32[4,128] { + %param_0.1203 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) + %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} + %constant.1221 = f32[]{:T(128)} constant(0) + %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1221), dimensions={}, metadata={op_name="broadcast.99"} + ROOT %mul.1920 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1203, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} } %fused_computation.431 () -> f32[64] { @@ -1425,35 +1425,35 @@ StackFrames ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.840, %div.768), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} } -%fused_computation.432 (param_0.1218: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.1218 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1218), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} +%fused_computation.432 (param_0.1217: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { + %param_0.1217 = s32[4,128]{1,0:T(4,128)} parameter(0) + %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1217), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %bitcast.418 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} ROOT %tuple.162 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.418, %convert_element_type.1405) } -%fused_computation.435 (param_0.1360: f32[2048,4]) -> bf16[4,2048] { - %param_0.1360 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.531 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.531) +%fused_computation.435 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { + %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) + %bitcast.533 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.533) } -%fused_computation.436 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { - %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.530 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.530) +%fused_computation.436 (param_0.1358: f32[2048,4]) -> bf16[4,2048] { + %param_0.1358 = f32[2048,4]{0,1:T(4,128)} parameter(0) + %bitcast.532 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.532) } -%fused_computation.437 (param_0.1361: f32[128,4]) -> bf16[4,128] { - %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.532 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.532) +%fused_computation.437 (param_0.1360: f32[128,4]) -> bf16[4,128] { + %param_0.1360 = f32[128,4]{0,1:T(4,128)} parameter(0) + %bitcast.534 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.534) } -%fused_computation.438 (param_0.1362: f32[128,4]) -> bf16[4,128] { - %param_0.1362 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.533 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1362), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.533) +%fused_computation.438 (param_0.1361: f32[128,4]) -> bf16[4,128] { + %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) + %bitcast.535 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.535) } %region_8.11 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { @@ -1462,40 +1462,40 @@ StackFrames ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.287.clone.clone (param_0.1346: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1346 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.526 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1346), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%fused_computation.287.clone.clone (param_0.1345: bf16[151936,2048]) -> bf16[151936,2048,1] { + %param_0.1345 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1345), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} } -%fused_computation.368.clone.clone (param_0.1347: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1281: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1281 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.476 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1281), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1438 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1347 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2067 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1347), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2066 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1438, %mul.2067), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2066), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.476, %convert_element_type.1437), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.368.clone.clone (param_0.1346: f32[4,128], param_1.1535: bf16[4,128,2048], param_2.1278: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1278 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1278), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1535 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1432 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1535), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1346 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2059 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1346), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.2058 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1432, %mul.2059), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1431 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2058), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.474 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.475, %convert_element_type.1431), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.439 (param_0.1363: bf16[151936,2048], param_1.1551: f32[4,128], param_2.1305: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { - %param_1.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1305 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) +%fused_computation.439 (param_0.1362: bf16[151936,2048], param_1.1544: f32[4,128], param_2.1302: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { + %param_1.1544 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1302 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) %param_3.913 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.270.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1551, %param_2.1305, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1363 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %fusion.253.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1363), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.270.clone.1, %fusion.253.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %constant.1195 = bf16[]{:T(256)} constant(-inf) - %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1195), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1544, %param_2.1302, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1362 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %fusion.252.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1362), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.269.clone.1, %fusion.252.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %constant.1193 = bf16[]{:T(256)} constant(-inf) + %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1193), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} ROOT %tuple.164 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,151936]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.223, %convolution.85.clone.1) } -%fused_computation.440 (param_0.1358: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.1358 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %bitcast.529 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.529) +%fused_computation.440 (param_0.1357: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { + %param_0.1357 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %bitcast.531 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1357), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.531) } %convert_element_type.767.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { @@ -1504,13 +1504,13 @@ StackFrames ROOT %add.755 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.155.clone.clone (param_0.1534: bf16[4,2048], param_1.1687: s32[]) -> bf16[2048] { - %param_0.1534 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) - %constant.1361 = s32[]{:T(128)} constant(0) - %dynamic_slice.388 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1687, %constant.1361), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1362 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.388, %constant.1362), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.155.clone.clone (param_0.1533: bf16[4,2048], param_1.1680: s32[]) -> bf16[2048] { + %param_0.1533 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1680 = s32[]{:T(128)S(6)} parameter(1) + %constant.1359 = s32[]{:T(128)} constant(0) + %dynamic_slice.384 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1533, %param_1.1680, %constant.1359), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1360 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.384, %constant.1360), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %region_14.16 (reduce_sum.204: f32[], reduce_sum.205: f32[]) -> f32[] { @@ -1519,25 +1519,25 @@ StackFrames ROOT %reduce_sum.206 = f32[]{:T(128)} add(%reduce_sum.204, %reduce_sum.205), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.58.clone.clone (param_0.1535: bf16[4,4,128,2048], param_1.1688: s32[]) -> f32[4,128] { - %param_0.1535 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1688 = s32[]{:T(128)S(6)} parameter(1) - %constant.1363 = s32[]{:T(128)} constant(0) - %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1535, %param_1.1688, %constant.1363, %constant.1363, %constant.1363), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1564 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.633), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.280 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1564, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1364 = f32[]{:T(128)} constant(0) - ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.280, %constant.1364), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} -} - -%fused_computation.179.clone.1.clone (param_0.1536: f32[4,128]) -> f32[4,128] { - %param_0.1536 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1366 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1366), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1536, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1365 = f32[]{:T(128)} constant(1e-06) - %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1365), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} +%fused_computation.58.clone.clone (param_0.1534: bf16[4,4,128,2048], param_1.1681: s32[]) -> f32[4,128] { + %param_0.1534 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1681 = s32[]{:T(128)S(6)} parameter(1) + %constant.1361 = s32[]{:T(128)} constant(0) + %dynamic_slice.385 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1681, %constant.1361, %constant.1361, %constant.1361), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.635 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1558 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.635), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.280 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1558, %convert_element_type.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1362 = f32[]{:T(128)} constant(0) + ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.280, %constant.1362), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} +} + +%fused_computation.179.clone.1.clone (param_0.1535: f32[4,128]) -> f32[4,128] { + %param_0.1535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %constant.1364 = f32[]{:T(128)} constant(0.00048828125) + %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1364), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1535, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1363 = f32[]{:T(128)} constant(1e-06) + %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1363), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.1039 = f32[4,128]{1,0:T(4,128)} add(%div.999, %closed_call.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.181 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1039), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } @@ -1548,158 +1548,158 @@ StackFrames ROOT %reduce_sum.212 = f32[]{:T(128)} add(%reduce_sum.207, %reduce_sum.211), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1550: bf16[4,2048,16,128], param_1.1698: s32[]) -> bf16[2048,16,128,1] { - %param_0.1550 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1698 = s32[]{:T(128)S(6)} parameter(1) - %constant.1377 = s32[]{:T(128)} constant(0) - %dynamic_slice.395 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1550, %param_1.1698, %constant.1377, %constant.1377, %constant.1377), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.644 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.114.clone.clone.clone.clone (param_0.1551: f32[4,128], param_1.1699: bf16[4,4,128,2048], param_2.1405: s32[], param_3.982: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.982 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.571 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1699 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1405 = s32[]{:T(128)S(6)} parameter(2) - %constant.1378 = s32[]{:T(128)} constant(0) - %dynamic_slice.396 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1699, %param_2.1405, %constant.1378, %constant.1378, %constant.1378), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.646 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.646), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2256 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1551), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2255 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %mul.2256), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2255), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.571, %convert_element_type.1574), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.645 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.61.clone.clone (param_0.1552: bf16[4,2048,16,128], param_1.1700: s32[], param_2.1406: f32[4,128], param_3.983: bf16[4,4,128,2048], param_4.604: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { - %param_2.1406 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.983 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1700 = s32[]{:T(128)S(6)} parameter(1) - %param_4.604 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1406, %param_3.983, %param_1.1700, %param_4.604), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1552 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1552, %param_1.1700), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.44.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1576 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.44.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.282 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1576, %convert_element_type.1576), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1379 = f32[]{:T(128)} constant(0) - %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.282, %constant.1379), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.44.clone.3) -} - -%fused_computation.151.clone.1.clone (param_0.1553: f32[4,128,16]) -> f32[4,128,16] { - %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1380 = f32[]{:T(128)} constant(0.0078125) - %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1380), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1553, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1381 = f32[]{:T(128)} constant(1e-06) - %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} +%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1549: bf16[4,2048,16,128], param_1.1691: s32[]) -> bf16[2048,16,128,1] { + %param_0.1549 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1691 = s32[]{:T(128)S(6)} parameter(1) + %constant.1375 = s32[]{:T(128)} constant(0) + %dynamic_slice.391 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1691, %constant.1375, %constant.1375, %constant.1375), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.646 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.114.clone.clone.clone.clone (param_0.1550: f32[4,128], param_1.1692: bf16[4,4,128,2048], param_2.1403: s32[], param_3.983: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.983 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.983), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1692 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1403 = s32[]{:T(128)S(6)} parameter(2) + %constant.1376 = s32[]{:T(128)} constant(0) + %dynamic_slice.392 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1692, %param_2.1403, %constant.1376, %constant.1376, %constant.1376), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.648 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1569 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.648), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1550 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2248 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2247 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1569, %mul.2248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1568 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2247), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.569 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.570, %convert_element_type.1568), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.647 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.61.clone.clone (param_0.1551: bf16[4,2048,16,128], param_1.1693: s32[], param_2.1404: f32[4,128], param_3.984: bf16[4,4,128,2048], param_4.603: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { + %param_2.1404 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.984 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1693 = s32[]{:T(128)S(6)} parameter(1) + %param_4.603 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1404, %param_3.984, %param_1.1693, %param_4.603), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1551 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1551, %param_1.1693), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.46.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %convert_element_type.1570 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.46.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.282 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1570, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1377 = f32[]{:T(128)} constant(0) + %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.282, %constant.1377), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.46.clone.3) +} + +%fused_computation.151.clone.1.clone (param_0.1552: f32[4,128,16]) -> f32[4,128,16] { + %param_0.1552 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) + %constant.1378 = f32[]{:T(128)} constant(0.0078125) + %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1378), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1552, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1379 = f32[]{:T(128)} constant(1e-06) + %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1379), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} %add.1043 = f32[4,128,16]{1,2,0:T(8,128)} add(%div.1001, %add.1044), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.183 = f32[4,128,16]{1,2,0:T(8,128)S(1)} rsqrt(%add.1043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.182.clone.clone (param_0.1549: bf16[4,128], param_1.1697: s32[]) -> bf16[128] { - %param_0.1549 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) - %constant.1376 = s32[]{:T(128)} constant(0) - %dynamic_slice.394 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1697, %constant.1376), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.643 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.394), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.182.clone.clone (param_0.1548: bf16[4,128], param_1.1690: s32[]) -> bf16[128] { + %param_0.1548 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1690 = s32[]{:T(128)S(6)} parameter(1) + %constant.1374 = s32[]{:T(128)} constant(0) + %dynamic_slice.390 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1548, %param_1.1690, %constant.1374), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.645 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.121.clone.1.clone (param_0.1554: f32[4,128,16], param_1.1701: bf16[4,128,16,128], param_2.1407: bf16[128]) -> bf16[4,128,16,128] { - %param_2.1407 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1407), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1701 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1578 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1554 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2258 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1554), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2257 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1578, %mul.2258), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1577 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2257), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.573, %convert_element_type.1577), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.121.clone.1.clone (param_0.1553: f32[4,128,16], param_1.1694: bf16[4,128,16,128], param_2.1405: bf16[128]) -> bf16[4,128,16,128] { + %param_2.1405 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1405), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1694 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1572 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1694), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) + %mul.2250 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1553), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2249 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1572, %mul.2250), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1571 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.571 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.572, %convert_element_type.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.90.clone.clone (param_0.1555: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { - %param_0.1555 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.90.clone.clone (param_0.1554: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { + %param_0.1554 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1554), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.129 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1554), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.209 = (bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) } %fused_computation.187.clone.clone () -> f32[64] { - %constant.1355 = f32[]{:T(128)} constant(1e+06) - %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1355), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %constant.1353 = f32[]{:T(128)} constant(1e+06) + %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1353), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} - %constant.1354 = s32[]{:T(128)} constant(2) - %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2242 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1562 = f32[64]{0:T(128)} convert(%mul.2242), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %constant.1356 = f32[]{:T(128)} constant(0.0078125) - %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1356), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1562, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1352 = s32[]{:T(128)} constant(2) + %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1352), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %mul.2234 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1556 = f32[64]{0:T(128)} convert(%mul.2234), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %constant.1354 = f32[]{:T(128)} constant(0.0078125) + %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1556, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.104, %div.995), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} } -%fused_computation.143.clone.clone (param_0.1529: f32[64], param_1.1683: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1683 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1683), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1529 = f32[64]{0:T(128)S(1)} parameter(0) - %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1529), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} +%fused_computation.143.clone.clone (param_0.1528: f32[64], param_1.1676: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1676 = f32[4,128]{1,0:T(4,128)} parameter(1) + %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1676), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %param_0.1528 = f32[64]{0:T(128)S(1)} parameter(0) + %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1528), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %div.996 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.998, %div.997), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1563 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1557 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} %convert_element_type.1189.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1563, %convert_element_type.1189.clone.3) + ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1557, %convert_element_type.1189.clone.3) } -%fused_computation.146.clone.1.clone (param_0.1530: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1530 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1357 = bf16[]{:T(256)} constant(-inf) - %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} +%fused_computation.146.clone.1.clone (param_0.1529: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1529 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) + %constant.1355 = bf16[]{:T(256)} constant(-inf) + %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1529, %constant.1355), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1529, %constant.1355), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.53 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.69, %pad.68), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.630 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.632 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.145.clone.1.clone (param_0.1545: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1545 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1374 = bf16[]{:T(256)} constant(-inf) - %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} +%fused_computation.145.clone.1.clone (param_0.1544: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1544 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) + %constant.1372 = bf16[]{:T(256)} constant(-inf) + %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1544, %constant.1372), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1544, %constant.1372), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.54 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.71, %pad.70), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.641 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.643 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.94.clone.clone (param_0.1556: bf16[4,128,16,64], param_1.1702: bf16[4,128,16,64], param_2.1408: bf16[4,128,128], param_3.984: bf16[4,128,128], param_4.605: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.384: bf16[128]) -> bf16[4,16,128,128] { - %param_6.384 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.575 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.384), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.94.clone.clone (param_0.1555: bf16[4,128,16,64], param_1.1695: bf16[4,128,16,64], param_2.1406: bf16[4,128,128], param_3.985: bf16[4,128,128], param_4.604: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.383: bf16[128]) -> bf16[4,16,128,128] { + %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} %param_5.499 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1580 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.605 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2265 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.605), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2264 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1580, %mul.2265), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1579 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.575, %convert_element_type.1579), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.984 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2263 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.984), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2261 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %mul.2263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1702 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1382 = bf16[]{:T(256)} constant(-inf) - %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1702, %constant.1382), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1556 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1556, %constant.1382), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %convert_element_type.1574 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_4.604 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) + %mul.2257 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.604), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2256 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1574, %mul.2257), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2256), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %convert_element_type.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.985 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2255 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.985), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2253 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.573, %mul.2255), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1695 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1380 = bf16[]{:T(256)} constant(-inf) + %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1695, %constant.1380), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1555 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1555, %constant.1380), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.56 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.75, %pad.74), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1408 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2262 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1408), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2260 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2262), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2261, %mul.2260), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %constant.1383 = bf16[]{:T(256)} constant(0.08838) - %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1383), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2259 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - ROOT %bitcast.647 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2259), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_2.1406 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2254 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1406), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2252 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2254), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2253, %mul.2252), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + %constant.1381 = bf16[]{:T(256)} constant(0.08838) + %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %mul.2251 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.649 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } %region_16.18 (reduce_sum.213: f32[], reduce_sum.214: f32[]) -> f32[] { @@ -1708,159 +1708,159 @@ StackFrames ROOT %reduce_sum.218 = f32[]{:T(128)} add(%reduce_sum.213, %reduce_sum.214), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1541: bf16[4,2048,8,128], param_1.1692: s32[]) -> bf16[2048,8,128,1] { - %param_0.1541 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1692 = s32[]{:T(128)S(6)} parameter(1) - %constant.1369 = s32[]{:T(128)} constant(0) - %dynamic_slice.392 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1541, %param_1.1692, %constant.1369, %constant.1369, %constant.1369), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.638 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.113.clone.clone.clone.clone (param_0.1542: f32[4,128], param_1.1693: bf16[4,4,128,2048], param_2.1401: s32[], param_3.979: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.979 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.979), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1693 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1401 = s32[]{:T(128)S(6)} parameter(2) - %constant.1370 = s32[]{:T(128)} constant(0) - %dynamic_slice.393 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1693, %param_2.1401, %constant.1370, %constant.1370, %constant.1370), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.640 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1568 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.640), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1542 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2246 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1542), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2245 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1568, %mul.2246), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.565, %convert_element_type.1567), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.639 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.84.clone.clone (param_0.1543: bf16[4,2048,8,128], param_1.1694: s32[], param_2.1402: f32[4,128], param_3.980: bf16[4,4,128,2048], param_4.602: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { - %param_2.1402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.980 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1694 = s32[]{:T(128)S(6)} parameter(1) - %param_4.602 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1402, %param_3.980, %param_1.1694, %param_4.602), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1543 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1543, %param_1.1694), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1540: bf16[4,2048,8,128], param_1.1685: s32[]) -> bf16[2048,8,128,1] { + %param_0.1540 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) + %constant.1367 = s32[]{:T(128)} constant(0) + %dynamic_slice.388 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1540, %param_1.1685, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.640 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.113.clone.clone.clone.clone (param_0.1541: f32[4,128], param_1.1686: bf16[4,4,128,2048], param_2.1399: s32[], param_3.980: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.980 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.980), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1686 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) + %constant.1368 = s32[]{:T(128)} constant(0) + %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1686, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.642 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1562 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.642), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1541 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2238 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1541), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2237 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1562, %mul.2238), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1561 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2237), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.564, %convert_element_type.1561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.641 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.84.clone.clone (param_0.1542: bf16[4,2048,8,128], param_1.1687: s32[], param_2.1400: f32[4,128], param_3.981: bf16[4,4,128,2048], param_4.601: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { + %param_2.1400 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.981 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) + %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1400, %param_3.981, %param_1.1687, %param_4.601), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1542 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1542, %param_1.1687), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %convolution.50.clone.3 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.73.clone.3, %fusion.87.clone.3), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1569 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.281 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1569, %convert_element_type.1569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1371 = f32[]{:T(128)} constant(0) - %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.281, %constant.1371), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %convert_element_type.1563 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.281 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1563, %convert_element_type.1563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1369 = f32[]{:T(128)} constant(0) + %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.281, %constant.1369), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.206 = (f32[4,128,8]{1,2,0:T(8,128)S(1)}, bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.246, %convolution.50.clone.3) } -%fused_computation.154.clone.1.clone (param_0.1544: f32[4,128,8]) -> f32[4,128,8] { - %param_0.1544 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1372 = f32[]{:T(128)} constant(0.0078125) - %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1372), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1544, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1373 = f32[]{:T(128)} constant(1e-06) - %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1373), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} +%fused_computation.154.clone.1.clone (param_0.1543: f32[4,128,8]) -> f32[4,128,8] { + %param_0.1543 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) + %constant.1370 = f32[]{:T(128)} constant(0.0078125) + %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1370), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1543, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1371 = f32[]{:T(128)} constant(1e-06) + %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1371), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} %add.1040 = f32[4,128,8]{1,2,0:T(8,128)} add(%div.1000, %add.1041), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.182 = f32[4,128,8]{1,2,0:T(8,128)S(1)} rsqrt(%add.1040), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.184.clone.clone (param_0.1528: bf16[4,128], param_1.1682: s32[]) -> bf16[128] { - %param_0.1528 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) - %constant.1353 = s32[]{:T(128)} constant(0) - %dynamic_slice.385 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1528, %param_1.1682, %constant.1353), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.629 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.139.clone.1.clone (param_0.1546: f32[4,128,8], param_1.1695: bf16[4,128,8,128], param_2.1403: bf16[128]) -> bf16[4,128,8,128] { - %param_2.1403 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1403), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1695 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1571 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1695), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1546 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2248 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1546), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2247 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1571, %mul.2248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1570 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2247), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.567, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.126.clone.clone (param_0.1547: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1547 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - ROOT %tuple.207 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) +%fused_computation.184.clone.clone (param_0.1527: bf16[4,128], param_1.1675: s32[]) -> bf16[128] { + %param_0.1527 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1675 = s32[]{:T(128)S(6)} parameter(1) + %constant.1351 = s32[]{:T(128)} constant(0) + %dynamic_slice.381 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1527, %param_1.1675, %constant.1351), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.631 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.129.clone.clone (param_0.1548: bf16[4,128,8,64], param_1.1696: bf16[4,128,8,64], param_2.1404: bf16[4,128,128], param_3.981: bf16[4,128,128], param_4.603: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.383: bf16[128]) -> bf16[4,8,128,128] { - %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.569 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_5.498 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1573 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.603 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2254 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.603), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2253 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1573, %mul.2254), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1572 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2253), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.569, %convert_element_type.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.981 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2252 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.981), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2250 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %mul.2252), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1696 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1375 = bf16[]{:T(256)} constant(-inf) - %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1696, %constant.1375), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1548 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1548, %constant.1375), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.55 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.73, %pad.72), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1404 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2251 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1404), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2249 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2250, %mul.2249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.642 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.169.clone.clone (param_0.1537: bf16[4,2048,8,128], param_1.1689: s32[]) -> bf16[1,2048,8,128] { - %param_0.1537 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1689 = s32[]{:T(128)S(6)} parameter(1) - %constant.1367 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.390 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1537, %param_1.1689, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +%fused_computation.139.clone.1.clone (param_0.1545: f32[4,128,8], param_1.1688: bf16[4,128,8,128], param_2.1401: bf16[128]) -> bf16[4,128,8,128] { + %param_2.1401 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1401), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1688 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1565 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1688), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1545 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) + %mul.2240 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1545), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2239 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1565, %mul.2240), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1564 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2239), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.565 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.566, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1538: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { - %param_0.1538 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.204 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1538), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.634 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.204), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.126.clone.clone (param_0.1546: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { + %param_0.1546 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1546), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} + %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1546), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + ROOT %tuple.207 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) } -%fused_computation.111.clone.clone.clone.clone (param_0.1539: f32[4,128], param_1.1690: bf16[4,4,128,2048], param_2.1399: s32[], param_3.977: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.977 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.977), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1690 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) - %constant.1368 = s32[]{:T(128)} constant(0) - %dynamic_slice.391 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1690, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.636 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1566 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.636), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1539 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2244 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2243 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1566, %mul.2244), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2243), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.563, %convert_element_type.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.635 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.562), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.140.clone.clone (param_0.1540: bf16[1,2048,8,128], param_1.1691: f32[4,128], param_2.1400: bf16[4,4,128,2048], param_3.978: s32[], param_4.601: bf16[2048]) -> bf16[4,8,128,128] { - %param_1.1691 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1400 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.978 = s32[]{:T(128)S(6)} parameter(3) - %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.373 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1691, %param_2.1400, %param_3.978, %param_4.601), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1540 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.372 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1540), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.373, %fusion.372), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.637 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.129.clone.clone (param_0.1547: bf16[4,128,8,64], param_1.1689: bf16[4,128,8,64], param_2.1402: bf16[4,128,128], param_3.982: bf16[4,128,128], param_4.602: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.382: bf16[128]) -> bf16[4,8,128,128] { + %param_6.382 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.382), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_5.498 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) + %convert_element_type.1567 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_4.602 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) + %mul.2246 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.602), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2245 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1567, %mul.2246), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %convert_element_type.1566), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.982 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2244 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2242 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.567, %mul.2244), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1689 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1373 = bf16[]{:T(256)} constant(-inf) + %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1689, %constant.1373), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1547 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1547, %constant.1373), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %maximum.55 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.73, %pad.72), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_2.1402 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2243 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1402), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2241 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2243), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2242, %mul.2241), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.644 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.188.clone.clone (param_0.1578: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { - %param_0.1578 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1578), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.660 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.660, %slice.11) +%fused_computation.169.clone.clone (param_0.1536: bf16[4,2048,8,128], param_1.1682: s32[]) -> bf16[1,2048,8,128] { + %param_0.1536 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) + %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) + %constant.1365 = s32[]{:T(128)} constant(0) + ROOT %dynamic_slice.386 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1536, %param_1.1682, %constant.1365, %constant.1365, %constant.1365), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1537: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { + %param_0.1537 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %copy.208 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1537), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} + ROOT %bitcast.636 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.208), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.111.clone.clone.clone.clone (param_0.1538: f32[4,128], param_1.1683: bf16[4,4,128,2048], param_2.1397: s32[], param_3.978: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.978 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.978), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1683 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1397 = s32[]{:T(128)S(6)} parameter(2) + %constant.1366 = s32[]{:T(128)} constant(0) + %dynamic_slice.387 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1683, %param_2.1397, %constant.1366, %constant.1366, %constant.1366), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.638 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1560 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.638), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1538 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2236 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1538), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2235 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1560, %mul.2236), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1559 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2235), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.561 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.562, %convert_element_type.1559), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.637 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.140.clone.clone (param_0.1539: bf16[1,2048,8,128], param_1.1684: f32[4,128], param_2.1398: bf16[4,4,128,2048], param_3.979: s32[], param_4.600: bf16[2048]) -> bf16[4,8,128,128] { + %param_1.1684 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1398 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.979 = s32[]{:T(128)S(6)} parameter(3) + %param_4.600 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.372 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1684, %param_2.1398, %param_3.979, %param_4.600), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1539 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.371 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1539), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.372, %fusion.371), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + ROOT %bitcast.639 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.188.clone.clone (param_0.1577: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { + %param_0.1577 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) + %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1577), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.662 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} + ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.662, %slice.11) } %region_17.20 (reduce_sum.219: f32[], reduce_sum.220: f32[]) -> f32[] { @@ -1869,36 +1869,36 @@ StackFrames ROOT %reduce_sum.221 = f32[]{:T(128)} add(%reduce_sum.219, %reduce_sum.220), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,2048], param_1.1703: s32[]) -> bf16[16,128,2048,1] { - %param_0.1557 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) - %constant.1384 = s32[]{:T(128)} constant(0) - %dynamic_slice.397 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1557, %param_1.1703, %constant.1384, %constant.1384, %constant.1384), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.648 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.397), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1556: bf16[4,16,128,2048], param_1.1696: s32[]) -> bf16[16,128,2048,1] { + %param_0.1556 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1696 = s32[]{:T(128)S(6)} parameter(1) + %constant.1382 = s32[]{:T(128)} constant(0) + %dynamic_slice.393 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1556, %param_1.1696, %constant.1382, %constant.1382, %constant.1382), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.650 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1558: bf16[4,16,128,128]) -> bf16[4,128,16,128] { - %param_0.1558 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.649 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,128]) -> bf16[4,128,16,128] { + %param_0.1557 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.651 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.64.clone.clone (param_0.1559: bf16[4,16,128,2048], param_1.1704: s32[], param_2.1409: bf16[4,16,128,128], param_3.985: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { - %param_3.985 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1704 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.64.clone.clone (param_0.1558: bf16[4,16,128,2048], param_1.1697: s32[], param_2.1407: bf16[4,16,128,128], param_3.986: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { + %param_3.986 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) %constant.436.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.242.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.985, %param_1.1704, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.242.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1409 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1409), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1559 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1559, %param_1.1704), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.240.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.986, %param_1.1697, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.240.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %param_2.1407 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1407), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_0.1558 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1558, %param_1.1697), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %convolution.62.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.96.clone.3, %fusion.95.clone.3), window={size=1x16}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %bitcast.203.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %add.768.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.227.clone.3, %bitcast.203.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1581 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.283 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1581, %convert_element_type.1581), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1385 = f32[]{:T(128)} constant(0) - %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.283, %constant.1385), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.283 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %convert_element_type.1575), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1383 = f32[]{:T(128)} constant(0) + %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.283, %constant.1383), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.210 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.248, %add.768.clone.3) } @@ -1908,93 +1908,93 @@ StackFrames ROOT %add.754 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.156.clone.clone (param_0.1531: bf16[4,2048], param_1.1684: s32[]) -> bf16[2048] { - %param_0.1531 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1684 = s32[]{:T(128)S(6)} parameter(1) - %constant.1358 = s32[]{:T(128)} constant(0) - %dynamic_slice.386 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1684, %constant.1358), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1359 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.386, %constant.1359), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.156.clone.clone (param_0.1530: bf16[4,2048], param_1.1677: s32[]) -> bf16[2048] { + %param_0.1530 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1677 = s32[]{:T(128)S(6)} parameter(1) + %constant.1356 = s32[]{:T(128)} constant(0) + %dynamic_slice.382 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1530, %param_1.1677, %constant.1356), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1357 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.382, %constant.1357), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.13.clone.clone.clone (param_0.1532: bf16[4,6144,2048], param_1.1685: s32[]) -> bf16[6144,2048,1] { - %param_0.1532 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) - %constant.1360 = s32[]{:T(128)} constant(0) - %dynamic_slice.387 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1532, %param_1.1685, %constant.1360, %constant.1360), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.632 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.13.clone.clone.clone (param_0.1531: bf16[4,6144,2048], param_1.1678: s32[]) -> bf16[6144,2048,1] { + %param_0.1531 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1678 = s32[]{:T(128)S(6)} parameter(1) + %constant.1358 = s32[]{:T(128)} constant(0) + %dynamic_slice.383 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1678, %constant.1358, %constant.1358), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.634 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %bitcast_fusion.1.clone.clone (bitcast_input.4: bf16[4,128,2048]) -> bf16[4,128,2048] { - %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.631 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) -} - -%fused_computation.14.clone.clone (param_0.1533: bf16[4,128,2048], param_1.1686: bf16[4,6144,2048], param_2.1398: s32[]) -> bf16[6144,4,128] { - %param_1.1686 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1398 = s32[]{:T(128)S(6)} parameter(2) - %fusion.370 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1686, %param_2.1398), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1533 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %fusion.371 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1533), kind=kLoop, calls=%bitcast_fusion.1.clone.clone - ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.370, %fusion.371), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.180.clone.1.clone (param_0.1560: f32[4,128]) -> f32[4,128] { - %param_0.1560 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1387 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1387), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1560, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1386 = f32[]{:T(128)} constant(1e-06) - %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1386), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) +} + +%fused_computation.14.clone.clone (param_0.1532: bf16[4,128,2048], param_1.1679: bf16[4,6144,2048], param_2.1396: s32[]) -> bf16[6144,4,128] { + %param_1.1679 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1396 = s32[]{:T(128)S(6)} parameter(2) + %fusion.369 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1679, %param_2.1396), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1532 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.370 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1532), kind=kLoop, calls=%bitcast_fusion.1.clone.clone + ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.369, %fusion.370), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.180.clone.1.clone (param_0.1559: f32[4,128]) -> f32[4,128] { + %param_0.1559 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %constant.1385 = f32[]{:T(128)} constant(0.00048828125) + %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1385), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1559, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1384 = f32[]{:T(128)} constant(1e-06) + %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1384), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.1046 = f32[4,128]{1,0:T(4,128)} add(%div.1002, %closed_call.110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.184 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1046), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.12.clone.1.clone.clone (param_0.1564: bf16[4,2048,6144], param_1.1708: s32[]) -> bf16[2048,6144,1] { - %param_0.1564 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1708 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.12.clone.1.clone.clone (param_0.1563: bf16[4,2048,6144], param_1.1701: s32[]) -> bf16[2048,6144,1] { + %param_0.1563 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1701 = s32[]{:T(128)S(6)} parameter(1) + %constant.1387 = s32[]{:T(128)} constant(0) + %dynamic_slice.395 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1563, %param_1.1701, %constant.1387, %constant.1387), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.119.clone.3.clone.clone (param_0.1564: f32[4,128], param_1.1702: bf16[4,128,2048], param_2.1410: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1410 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1410), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1702 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1579 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1564 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2261 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1564), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2260 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1579, %mul.2261), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2260), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.577 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.578, %convert_element_type.1578), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.21.clone.clone (param_0.1565: bf16[4,2048,6144], param_1.1703: s32[], param_2.1411: f32[4,128], param_3.988: bf16[4,128,2048], param_4.606: bf16[2048]) -> bf16[4,128,6144] { + %param_2.1411 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.988 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.606 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.376 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1411, %param_3.988, %param_4.606), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1565 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) + %fusion.375 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1565, %param_1.1703), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.376, %fusion.375), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.11.clone.1.clone.clone (param_0.1567: bf16[4,2048,6144], param_1.1705: s32[]) -> bf16[2048,6144,1] { + %param_0.1567 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1705 = s32[]{:T(128)S(6)} parameter(1) %constant.1389 = s32[]{:T(128)} constant(0) - %dynamic_slice.399 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1564, %param_1.1708, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.651 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.119.clone.3.clone.clone (param_0.1565: f32[4,128], param_1.1709: bf16[4,128,2048], param_2.1412: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1412 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.579 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1412), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1709 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1585 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1565 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2269 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1565), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2268 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1585, %mul.2269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1584 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.579, %convert_element_type.1584), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.21.clone.clone (param_0.1566: bf16[4,2048,6144], param_1.1710: s32[], param_2.1413: f32[4,128], param_3.987: bf16[4,128,2048], param_4.607: bf16[2048]) -> bf16[4,128,6144] { - %param_2.1413 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.987 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.607 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.377 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1413, %param_3.987, %param_4.607), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1566 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1710 = s32[]{:T(128)S(6)} parameter(1) - %fusion.376 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1566, %param_1.1710), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.377, %fusion.376), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.11.clone.1.clone.clone (param_0.1568: bf16[4,2048,6144], param_1.1712: s32[]) -> bf16[2048,6144,1] { - %param_0.1568 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1712 = s32[]{:T(128)S(6)} parameter(1) - %constant.1391 = s32[]{:T(128)} constant(0) - %dynamic_slice.400 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1568, %param_1.1712, %constant.1391, %constant.1391), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.400), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.47.clone.1.clone.clone (param_0.1567: bf16[6144,4,128], param_1.1711: bf16[4,128,6144]) -> bf16[4,128,6144] { - %param_1.1711 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1390 = bf16[]{:T(256)} constant(1) - %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1390), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1711), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %dynamic_slice.396 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1567, %param_1.1705, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.655 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.47.clone.1.clone.clone (param_0.1566: bf16[6144,4,128], param_1.1704: bf16[4,128,6144]) -> bf16[4,128,6144] { + %param_1.1704 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1388 = bf16[]{:T(256)} constant(1) + %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1388), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} + %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.69 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.1047 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.1003 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.1047), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.2271 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1711, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %mul.2263 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1704, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} From 55eeb67e4a96594eff9fc686e028ae657a640c40 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Wed, 10 Jun 2026 22:31:07 +0000 Subject: [PATCH 31/52] Fix quant_einsum to correctly pass kwargs to einsum callables --- src/maxtext/layers/moe.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 28c3dfe5e3..6aa75df88e 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -1916,11 +1916,9 @@ def get_einsum( if self.quant: - def quant_einsum(*args, **kwargs): # pylint: disable=unused-argument - # simply skip kwargs, since einsum doesn't support any kwargs - # like precision + def quant_einsum(*args, **kwargs): kw = {"dtype": self.dtype} - return self.quant.einsum(**kw)(*args) # pytype: disable=attribute-error + return self.quant.einsum(**kw)(*args, **kwargs) # pytype: disable=attribute-error einsum_op = quant_einsum else: From fdb4da7df698a6830ec8d672d86c8e0e3288be12 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 11 Jun 2026 05:04:08 +0000 Subject: [PATCH 32/52] Update reference HLO from CI artifact --- tests/utils/reference_hlo_llama3_8b.txt | 2086 +++++++++++----------- tests/utils/reference_hlo_qwen3_1.7b.txt | 1906 ++++++++++---------- 2 files changed, 1996 insertions(+), 1996 deletions(-) diff --git a/tests/utils/reference_hlo_llama3_8b.txt b/tests/utils/reference_hlo_llama3_8b.txt index 488affcd35..27c6529df2 100644 --- a/tests/utils/reference_hlo_llama3_8b.txt +++ b/tests/utils/reference_hlo_llama3_8b.txt @@ -44,62 +44,62 @@ StackFrames ROOT %reduce_sum.192 = f32[]{:T(128)} add(%reduce_sum.190, %reduce_sum.191), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.281.clone.clone.clone (param_0.1085: bf16[4,128,128256], param_1.1251: s32[4,128], param_2.1077: f32[4,128], param_3.781: f32[4,128], param_4.482: bf16[4,128], param_5.404: f32[4,128]) -> bf16[4,128,128256] { - %param_5.404 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1607 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.404), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.781 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1606 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.781), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1085 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1032 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1085), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.482 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.482), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1032, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.280.clone.clone.clone (param_0.1099: bf16[4,128,128256], param_1.1265: s32[4,128], param_2.1086: f32[4,128], param_3.785: f32[4,128], param_4.487: bf16[4,128], param_5.412: f32[4,128]) -> bf16[4,128,128256] { + %param_5.412 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.1613 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.785 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.1612 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.785), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1099 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1044 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1099), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.487 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.487), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1044, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.62 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1605 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1606, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1077 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1077), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1605, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1251 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1251), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.1611 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1612, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1611, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1265 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1265), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.48 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.47 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1031 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1031), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1604 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1607, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1030 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1604), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.317.clone.clone (param_0.1086: f32[4,128], param_1.1252: bf16[4,128,4096], param_2.1079: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1079 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1079), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1034 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1252), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1609 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1608 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1034, %mul.1609), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1033 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1608), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.220 (param_0.1105: bf16[4,128,128256], param_1.1267: s32[4,128], param_2.1103: f32[4,128], param_3.797: f32[4,128], param_4.497: bf16[4,128], param_5.419: f32[4,128], param_6.287: f32[4,128], param_7.186: bf16[4,128,4096], param_8.112: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { - %param_6.287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.186 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) - %param_8.112 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) - %fusion.229.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.287, %param_7.186, %param_8.112), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1105 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1267 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1103 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.797 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %param_4.497 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %param_5.419 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1105, %param_1.1267, %param_2.1103, %param_3.797, %param_4.497, /*index=5*/%param_5.419), kind=kLoop, calls=%fused_computation.281.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.82.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.229.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.300 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.82.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.911 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.300), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %square.157 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.911, %convert_element_type.911), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %convert_element_type.1043 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1043), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1610 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1613, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1042 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1610), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.316.clone.clone (param_0.1100: f32[4,128], param_1.1266: bf16[4,128,4096], param_2.1088: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1088 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.387 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1088), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1266 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1046 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1266), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1100 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1615 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1100), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1614 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1046, %mul.1615), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1045 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1614), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.387, %convert_element_type.1045), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.219 (param_0.1119: bf16[4,128,128256], param_1.1281: s32[4,128], param_2.1112: f32[4,128], param_3.801: f32[4,128], param_4.502: bf16[4,128], param_5.427: f32[4,128], param_6.299: f32[4,128], param_7.198: bf16[4,128,4096], param_8.116: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { + %param_6.299 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.198 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) + %param_8.116 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) + %fusion.239.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.299, %param_7.198, %param_8.116), kind=kLoop, calls=%fused_computation.316.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1119 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1281 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1112 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.801 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %param_4.502 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %param_5.427 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1119, %param_1.1281, %param_2.1112, %param_3.801, %param_4.502, /*index=5*/%param_5.427), kind=kLoop, calls=%fused_computation.280.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convolution.88.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.239.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %bitcast.306 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.88.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.923 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.306), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %square.157 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.923, %convert_element_type.923), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1006 = f32[]{:T(128)} constant(0) %reduce.118 = f32[]{:T(128)} reduce(%square.157, %constant.1006), dimensions={0,1}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.82.clone.1) + ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.88.clone.1) } %region_34.39 (reduce_sum.196: f32[], reduce_sum.197: f32[]) -> f32[] { @@ -108,10 +108,10 @@ StackFrames ROOT %reduce_sum.198 = f32[]{:T(128)} add(%reduce_sum.196, %reduce_sum.197), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.221 (param_0.1104: bf16[128256,4096]) -> f32[] { - %param_0.1104 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.913 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1104), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %square.159 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.913, %convert_element_type.913), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.220 (param_0.1118: bf16[128256,4096]) -> f32[] { + %param_0.1118 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.925 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1118), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %square.159 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.925, %convert_element_type.925), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1005 = f32[]{:T(128)} constant(0) ROOT %reduce.119 = f32[]{:T(128)} reduce(%square.159, %constant.1005), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } @@ -128,39 +128,39 @@ StackFrames ROOT %reduce_sum.261 = f32[]{:T(128)} add(%reduce_sum.259, %reduce_sum.260), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.222 (param_0.1092: f32[128256,4096], param_1.1255: f32[], param_2.1091: f32[], param_3.785: f32[], param_4.485: f32[128256,4096], param_5.407: f32[], param_6.275: bf16[128256,4096], param_7.174: pred[], param_8.100: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { - %param_0.1092 = f32[128256,4096]{1,0:T(8,128)} parameter(0) - %param_3.785 = f32[]{:T(128)S(6)} parameter(3) - %mul.1476.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.785), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.174 = pred[]{:T(512)S(6)} parameter(7) - %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.174), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.275 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1005.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.275), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %param_5.407 = f32[]{:T(128)} parameter(5) - %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.407), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1005.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1005.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.221 (param_0.1106: f32[128256,4096], param_1.1269: f32[], param_2.1100: f32[], param_3.789: f32[], param_4.490: f32[128256,4096], param_5.415: f32[], param_6.287: bf16[128256,4096], param_7.186: pred[], param_8.104: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { + %param_0.1106 = f32[128256,4096]{1,0:T(8,128)} parameter(0) + %param_3.789 = f32[]{:T(128)S(6)} parameter(3) + %mul.1482.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.186 = pred[]{:T(512)S(6)} parameter(7) + %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.186), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.287 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1017.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.287), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_5.415 = f32[]{:T(128)} parameter(5) + %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1017.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1017.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.907.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.554.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.907.clone.1), dimensions={}, metadata={op_name="broadcast.61"} - %mul.1482.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.100 = f32[128256,4096]{1,0:T(8,128)} parameter(8) + %mul.1488.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.104 = f32[128256,4096]{1,0:T(8,128)} parameter(8) %constant.911.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1483.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1481.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.100, %mul.1483.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1482.clone.1, %mul.1481.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1091 = f32[]{:T(128)S(6)} parameter(2) - %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1091), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1489.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1487.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.104, %mul.1489.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1488.clone.1, %mul.1487.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) + %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.60.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %select_n.241.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.910.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1480.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1478.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1480.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.485 = f32[128256,4096]{1,0:T(8,128)} parameter(4) + %mul.1486.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1484.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1486.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.490 = f32[128256,4096]{1,0:T(8,128)} parameter(4) %constant.909.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1479.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1477.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.485, %mul.1479.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1478.clone.1, %mul.1477.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1255 = f32[]{:T(128)S(6)} parameter(1) - %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1255), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1485.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1483.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.490, %mul.1485.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1484.clone.1, %mul.1483.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1269 = f32[]{:T(128)S(6)} parameter(1) + %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1269), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.719.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.775.clone.1, %div.720.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.58.clone.1 = f32[128256,4096]{1,0:T(8,128)} sqrt(%div.719.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.908.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -168,10 +168,10 @@ StackFrames %add.773.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%sqrt.58.clone.1, %add.774.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.256.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%div.721.clone.1, %add.773.clone.1), metadata={op_name="multiply.42"} %div.718.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.776.clone.1, %multiply.256.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1475.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1092, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1475.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1474.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1476.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1092, %mul.1474.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1481.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1106, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1481.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1480.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1482.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1106, %mul.1480.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.160 = f32[128256,4096]{1,0:T(8,128)} multiply(%add.771.clone.1, %add.771.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.993 = f32[]{:T(128)} constant(0) %reduce.120 = f32[]{:T(128)} reduce(%square.160, %constant.993), dimensions={0,1}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -191,40 +191,40 @@ StackFrames ROOT %reduce_sum.255 = f32[]{:T(128)} add(%reduce_sum.253, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.223 (param_0.1093: f32[4096,128256], param_1.1256: f32[], param_2.1092: f32[], param_3.786: f32[], param_4.486: f32[4096,128256], param_5.408: f32[], param_6.276: bf16[4096,128256,1], param_7.175: pred[], param_8.101: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { - %param_0.1093 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %param_3.786 = f32[]{:T(128)S(6)} parameter(3) - %mul.1486.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.786), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.175 = pred[]{:T(512)S(6)} parameter(7) - %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.175), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.276 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.403.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.1007.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.403.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %param_5.408 = f32[]{:T(128)} parameter(5) - %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.408), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1007.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1007.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.222 (param_0.1107: f32[4096,128256], param_1.1270: f32[], param_2.1101: f32[], param_3.790: f32[], param_4.491: f32[4096,128256], param_5.416: f32[], param_6.288: bf16[4096,128256,1], param_7.187: pred[], param_8.105: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { + %param_0.1107 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %param_3.790 = f32[]{:T(128)S(6)} parameter(3) + %mul.1492.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.187 = pred[]{:T(512)S(6)} parameter(7) + %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.187), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.288 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) + %bitcast.409.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.288), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %convert_element_type.1019.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.409.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %param_5.416 = f32[]{:T(128)} parameter(5) + %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1019.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1019.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.913.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.556.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.913.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %mul.1492.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.101 = f32[4096,128256]{1,0:T(8,128)} parameter(8) + %mul.1498.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.105 = f32[4096,128256]{1,0:T(8,128)} parameter(8) %constant.917.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1493.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1491.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.101, %mul.1493.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1492.clone.1, %mul.1491.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1092 = f32[]{:T(128)S(6)} parameter(2) - %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1092), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1499.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1497.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.105, %mul.1499.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1498.clone.1, %mul.1497.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) + %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.61.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %select_n.245.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.916.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1490.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1488.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1490.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.486 = f32[4096,128256]{1,0:T(8,128)} parameter(4) + %mul.1496.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1494.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1496.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.491 = f32[4096,128256]{1,0:T(8,128)} parameter(4) %constant.915.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1489.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1487.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.486, %mul.1489.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1488.clone.1, %mul.1487.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1256 = f32[]{:T(128)S(6)} parameter(1) - %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1256), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1495.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1493.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.491, %mul.1495.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1494.clone.1, %mul.1493.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1270 = f32[]{:T(128)S(6)} parameter(1) + %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1270), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.727.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.781.clone.1, %div.728.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.59.clone.1 = f32[4096,128256]{1,0:T(8,128)} sqrt(%div.727.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.914.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -232,10 +232,10 @@ StackFrames %add.779.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%sqrt.59.clone.1, %add.780.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.257.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%div.729.clone.1, %add.779.clone.1), metadata={op_name="multiply.41"} %div.726.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.782.clone.1, %multiply.257.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1485.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1093, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1485.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1484.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1486.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1093, %mul.1484.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1491.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1107, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1491.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1490.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1492.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1107, %mul.1490.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.161 = f32[4096,128256]{1,0:T(8,128)} multiply(%add.777.clone.1, %add.777.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.994 = f32[]{:T(128)} constant(0) %reduce.121 = f32[]{:T(128)} reduce(%square.161, %constant.994), dimensions={0,1}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -249,10 +249,10 @@ StackFrames ROOT %reduce_sum.156 = f32[]{:T(128)} add(%reduce_sum.154, %reduce_sum.155), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.240 (param_0.1110: f32[4,14336,4096]) -> f32[] { - %param_0.1110 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) - %bitcast.308 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.164 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.308, %bitcast.308), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.239 (param_0.1124: f32[4,14336,4096]) -> f32[] { + %param_0.1124 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) + %bitcast.314 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1124), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.164 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.314, %bitcast.314), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1011 = f32[]{:T(128)} constant(0) ROOT %reduce.124 = f32[]{:T(128)} reduce(%square.164, %constant.1011), dimensions={0,1,2}, to_apply=%region_25.30, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } @@ -269,35 +269,35 @@ StackFrames ROOT %reduce_sum.147 = f32[]{:T(128)} add(%reduce_sum.142, %reduce_sum.143), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.242 (param_0.1111: f32[4,4096,14336], param_1.1270: f32[4,4096,14336]) -> (f32[], f32[]) { - %param_0.1111 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) - %bitcast.312 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.167 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.312, %bitcast.312), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.241 (param_0.1125: f32[4,4096,14336], param_1.1284: f32[4,4096,14336]) -> (f32[], f32[]) { + %param_0.1125 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) + %bitcast.318 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1125), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.167 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.318, %bitcast.318), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1012 = f32[]{:T(128)} constant(0) %reduce.125 = f32[]{:T(128)} reduce(%square.167, %constant.1012), dimensions={0,1,2}, to_apply=%region_24.29, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1270 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) - %bitcast.316.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1270), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.170.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.316.clone.1, %bitcast.316.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %param_1.1284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) + %bitcast.322.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.170.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.322.clone.1, %bitcast.322.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %reduce.126.clone.1 = f32[]{:T(128)} reduce(%square.170.clone.1, %constant.1012), dimensions={0,1,2}, to_apply=%region_23.28, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.155 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.125, %reduce.126.clone.1) } -%fused_computation.245 (param_0.681: f32[14336,4,4096]) -> bf16[4,14336,4096] { - %param_0.681 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.681), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.317 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.244 (param_0.694: f32[14336,4,4096]) -> bf16[4,14336,4096] { + %param_0.694 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) + %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.694), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} + ROOT %bitcast.323 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.246 (param_0.683: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.683 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.683), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.318 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.245 (param_0.696: f32[4096,4,14336]) -> bf16[4,4096,14336] { + %param_0.696 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.696), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} + ROOT %bitcast.324 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.247 (param_0.685: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.685 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.685), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.319 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.246 (param_0.698: f32[4096,4,14336]) -> bf16[4,4096,14336] { + %param_0.698 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.698), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} + ROOT %bitcast.325 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_52.57 (reduce_sum.289: f32[], reduce_sum.290: f32[]) -> f32[] { @@ -312,39 +312,39 @@ StackFrames ROOT %reduce_sum.219 = f32[]{:T(128)} add(%reduce_sum.217, %reduce_sum.218), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.248 (param_0.1100: f32[14336,4,4096], param_1.1263: f32[], param_2.1099: f32[], param_3.793: f32[], param_4.493: f32[14336,4,4096], param_5.415: f32[], param_6.283: f32[4,14336,4096], param_7.182: pred[], param_8.108: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { - %param_0.1100 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %param_3.793 = f32[]{:T(128)S(6)} parameter(3) - %mul.1544.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.182 = pred[]{:T(512)S(6)} parameter(7) - %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.182), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.283 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) - %bitcast.417.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.415 = f32[]{:T(128)} parameter(5) - %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.417.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.417.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.247 (param_0.1114: f32[14336,4,4096], param_1.1277: f32[], param_2.1108: f32[], param_3.797: f32[], param_4.498: f32[14336,4,4096], param_5.423: f32[], param_6.295: f32[4,14336,4096], param_7.194: pred[], param_8.112: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { + %param_0.1114 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) + %param_3.797 = f32[]{:T(128)S(6)} parameter(3) + %mul.1550.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.797), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.194 = pred[]{:T(512)S(6)} parameter(7) + %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.194), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.295 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) + %bitcast.423.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.295), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.423 = f32[]{:T(128)} parameter(5) + %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.423), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.423.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.423.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.955.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.586.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.955.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1550.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.108 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) + %mul.1556.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.112 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) %constant.959.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1551.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1549.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.108, %mul.1551.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1550.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1099 = f32[]{:T(128)S(6)} parameter(2) - %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1099), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1557.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1555.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.112, %mul.1557.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1556.clone.1, %mul.1555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1108 = f32[]{:T(128)S(6)} parameter(2) + %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1108), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.68.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %select_n.273.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.958.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1548.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1546.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1548.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.493 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) + %mul.1554.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1552.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.498 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) %constant.957.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1547.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1545.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.493, %mul.1547.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1546.clone.1, %mul.1545.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1263 = f32[]{:T(128)S(6)} parameter(1) - %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1263), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1553.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1551.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.498, %mul.1553.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1552.clone.1, %mul.1551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1277 = f32[]{:T(128)S(6)} parameter(1) + %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.783.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.819.clone.1, %div.784.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.66.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} sqrt(%div.783.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.956.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -352,10 +352,10 @@ StackFrames %add.817.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%sqrt.66.clone.1, %add.818.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.264.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%div.785.clone.1, %add.817.clone.1), metadata={op_name="multiply.34"} %div.782.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.820.clone.1, %multiply.264.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1543.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1100, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1543.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1542.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1544.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1100, %mul.1542.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1549.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1114, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1548.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1550.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1114, %mul.1548.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.171 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%add.815.clone.1, %add.815.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1001 = f32[]{:T(128)} constant(0) %reduce.127 = f32[]{:T(128)} reduce(%square.171, %constant.1001), dimensions={0,1,2}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -375,39 +375,39 @@ StackFrames ROOT %reduce_sum.213 = f32[]{:T(128)} add(%reduce_sum.211, %reduce_sum.212), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.249 (param_0.1101: f32[4096,4,14336], param_1.1264: f32[], param_2.1100: f32[], param_3.794: f32[], param_4.494: f32[4096,4,14336], param_5.416: f32[], param_6.284: f32[4,4096,14336], param_7.183: pred[], param_8.109: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1101 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.794 = f32[]{:T(128)S(6)} parameter(3) - %mul.1554.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.183 = pred[]{:T(512)S(6)} parameter(7) - %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.183), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.419.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.416 = f32[]{:T(128)} parameter(5) - %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.419.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.419.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.248 (param_0.1115: f32[4096,4,14336], param_1.1278: f32[], param_2.1109: f32[], param_3.798: f32[], param_4.499: f32[4096,4,14336], param_5.424: f32[], param_6.296: f32[4,4096,14336], param_7.195: pred[], param_8.113: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { + %param_0.1115 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.798 = f32[]{:T(128)S(6)} parameter(3) + %mul.1560.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.798), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.195 = pred[]{:T(512)S(6)} parameter(7) + %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.195), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.296 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.425.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.296), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.424 = f32[]{:T(128)} parameter(5) + %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.424), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.425.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.425.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.961.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.592.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.961.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1558.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.109 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) + %mul.1564.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.113 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) %constant.965.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.591.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.965.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1557.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.109, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1558.clone.1, %mul.1557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) - %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1563.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.113, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1564.clone.1, %mul.1563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1109 = f32[]{:T(128)S(6)} parameter(2) + %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1109), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.69.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %select_n.277.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.964.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.590.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.964.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1556.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.494 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) + %mul.1562.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.499 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) %constant.963.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.589.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.963.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1555.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.494, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1556.clone.1, %mul.1555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1264 = f32[]{:T(128)S(6)} parameter(1) - %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1264), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1561.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.499, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1562.clone.1, %mul.1561.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1278 = f32[]{:T(128)S(6)} parameter(1) + %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1278), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.791.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.824.clone.1, %div.792.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.67.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.791.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.962.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -415,10 +415,10 @@ StackFrames %add.823.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.67.clone.1, %broadcast.587.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.265.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.793.clone.1, %add.823.clone.1), metadata={op_name="multiply.33"} %div.790.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.825.clone.1, %multiply.265.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1553.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1101, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1553.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1552.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1554.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1101, %mul.1552.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1559.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1115, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1558.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1560.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1115, %mul.1558.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.172 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.821.clone.1, %add.821.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1002 = f32[]{:T(128)} constant(0) %reduce.128 = f32[]{:T(128)} reduce(%square.172, %constant.1002), dimensions={0,1,2}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -438,39 +438,39 @@ StackFrames ROOT %reduce_sum.210 = f32[]{:T(128)} add(%reduce_sum.205, %reduce_sum.206), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.250 (param_0.1102: f32[4096,4,14336], param_1.1265: f32[], param_2.1101: f32[], param_3.795: f32[], param_4.495: f32[4096,4,14336], param_5.417: f32[], param_6.285: f32[4,4096,14336], param_7.184: pred[], param_8.110: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1102 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.795 = f32[]{:T(128)S(6)} parameter(3) - %mul.1561.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.184 = pred[]{:T(512)S(6)} parameter(7) - %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.184), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.285 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.421.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.417 = f32[]{:T(128)} parameter(5) - %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.421.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.421.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.249 (param_0.1116: f32[4096,4,14336], param_1.1279: f32[], param_2.1110: f32[], param_3.799: f32[], param_4.500: f32[4096,4,14336], param_5.425: f32[], param_6.297: f32[4,4096,14336], param_7.196: pred[], param_8.114: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { + %param_0.1116 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.799 = f32[]{:T(128)S(6)} parameter(3) + %mul.1567.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.799), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.196 = pred[]{:T(512)S(6)} parameter(7) + %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.196), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.297 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.427.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.297), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.425 = f32[]{:T(128)} parameter(5) + %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.425), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.427.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.427.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.967.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.598.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.967.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1565.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.110 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) + %mul.1571.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.114 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) %constant.971.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.597.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.971.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1564.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.110, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1565.clone.1, %mul.1564.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) - %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1570.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.114, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1571.clone.1, %mul.1570.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1110 = f32[]{:T(128)S(6)} parameter(2) + %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1110), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.70.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %select_n.281.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.970.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.596.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.970.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1563.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.495 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) + %mul.1569.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.500 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) %constant.969.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.595.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.969.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1562.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.495, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1563.clone.1, %mul.1562.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1265 = f32[]{:T(128)S(6)} parameter(1) - %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1265), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1568.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.500, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1569.clone.1, %mul.1568.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1279 = f32[]{:T(128)S(6)} parameter(1) + %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1279), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.799.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.829.clone.1, %div.800.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.68.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.799.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.968.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -478,10 +478,10 @@ StackFrames %add.828.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.68.clone.1, %broadcast.593.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.266.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.801.clone.1, %add.828.clone.1), metadata={op_name="multiply.32"} %div.798.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.830.clone.1, %multiply.266.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1560.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1102, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1560.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1559.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1561.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1102, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1566.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1116, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1565.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1567.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1116, %mul.1565.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.173 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.826.clone.1, %add.826.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1003 = f32[]{:T(128)} constant(0) %reduce.129 = f32[]{:T(128)} reduce(%square.173, %constant.1003), dimensions={0,1,2}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -495,10 +495,10 @@ StackFrames ROOT %reduce_sum.183 = f32[]{:T(128)} add(%reduce_sum.178, %reduce_sum.182), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.268 (param_0.1106: f32[4,4096,32,128]) -> f32[] { - %param_0.1106 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.323 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.176 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.323, %bitcast.323), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.267 (param_0.1120: f32[4,4096,32,128]) -> f32[] { + %param_0.1120 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.329 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1120), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.176 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.329, %bitcast.329), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1007 = f32[]{:T(128)} constant(0) ROOT %reduce.133 = f32[]{:T(128)} reduce(%square.176, %constant.1007), dimensions={0,1,2,3}, to_apply=%region_30.35, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } @@ -509,18 +509,18 @@ StackFrames ROOT %reduce_sum.177 = f32[]{:T(128)} add(%reduce_sum.175, %reduce_sum.176), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.270 (param_0.1107: f32[4,32,128,4096]) -> f32[] { - %param_0.1107 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.327 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.179 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.327, %bitcast.327), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.269 (param_0.1121: f32[4,32,128,4096]) -> f32[] { + %param_0.1121 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.333 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1121), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.179 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.333, %bitcast.333), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1008 = f32[]{:T(128)} constant(0) ROOT %reduce.134 = f32[]{:T(128)} reduce(%square.179, %constant.1008), dimensions={0,1,2,3}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.271 (param_0.735: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { - %param_0.735 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.735), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.328 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.270 (param_0.748: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { + %param_0.748 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) + %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.748), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} + ROOT %bitcast.334 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_57.62 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { @@ -535,39 +535,39 @@ StackFrames ROOT %reduce_sum.246 = f32[]{:T(128)} add(%reduce_sum.241, %reduce_sum.245), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.272 (param_0.1095: f32[4096,4,32,128], param_1.1258: f32[], param_2.1094: f32[], param_3.788: f32[], param_4.488: f32[4096,4,32,128], param_5.410: f32[], param_6.278: f32[4,4096,32,128], param_7.177: pred[], param_8.103: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { - %param_0.1095 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.788 = f32[]{:T(128)S(6)} parameter(3) - %mul.1503.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.788), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.177 = pred[]{:T(512)S(6)} parameter(7) - %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.177), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.278 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.407.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.278), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.410 = f32[]{:T(128)} parameter(5) - %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.410), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.407.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.407.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.271 (param_0.1109: f32[4096,4,32,128], param_1.1272: f32[], param_2.1103: f32[], param_3.792: f32[], param_4.493: f32[4096,4,32,128], param_5.418: f32[], param_6.290: f32[4,4096,32,128], param_7.189: pred[], param_8.107: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { + %param_0.1109 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.792 = f32[]{:T(128)S(6)} parameter(3) + %mul.1509.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.189 = pred[]{:T(512)S(6)} parameter(7) + %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.189), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.290 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.413.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.290), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.418 = f32[]{:T(128)} parameter(5) + %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.413.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.413.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.925.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.564.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.925.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %mul.1509.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.103 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1515.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.107 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) %constant.929.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1510.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1508.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.103, %mul.1510.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1509.clone.1, %mul.1508.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1094 = f32[]{:T(128)S(6)} parameter(2) - %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1094), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1516.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1514.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.107, %mul.1516.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1515.clone.1, %mul.1514.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1103 = f32[]{:T(128)S(6)} parameter(2) + %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1103), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.63.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %select_n.253.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.928.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1507.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1505.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1507.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.488 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1513.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1511.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1513.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.493 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) %constant.927.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1506.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1504.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.488, %mul.1506.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1505.clone.1, %mul.1504.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1258 = f32[]{:T(128)S(6)} parameter(1) - %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1258), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1512.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1510.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.493, %mul.1512.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1511.clone.1, %mul.1510.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) + %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.743.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.792.clone.1, %div.744.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.61.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} sqrt(%div.743.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.926.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -575,10 +575,10 @@ StackFrames %add.790.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%sqrt.61.clone.1, %add.791.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.259.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%div.745.clone.1, %add.790.clone.1), metadata={op_name="multiply.39"} %div.742.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.793.clone.1, %multiply.259.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1502.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1095, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1502.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1501.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1503.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1095, %mul.1501.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1508.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1109, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1508.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1507.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1509.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1109, %mul.1507.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.180 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%add.788.clone.1, %add.788.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.996 = f32[]{:T(128)} constant(0) %reduce.135 = f32[]{:T(128)} reduce(%square.180, %constant.996), dimensions={0,1,2,3}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -598,39 +598,39 @@ StackFrames ROOT %reduce_sum.240 = f32[]{:T(128)} add(%reduce_sum.238, %reduce_sum.239), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.273 (param_0.1096: f32[32,4,128,4096], param_1.1259: f32[], param_2.1095: f32[], param_3.789: f32[], param_4.489: f32[32,4,128,4096], param_5.411: f32[], param_6.279: f32[4,32,128,4096], param_7.178: pred[], param_8.104: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { - %param_0.1096 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %param_3.789 = f32[]{:T(128)S(6)} parameter(3) - %mul.1513.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.178 = pred[]{:T(512)S(6)} parameter(7) - %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.178), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.279 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.409.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.411 = f32[]{:T(128)} parameter(5) - %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.411), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.409.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.409.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.272 (param_0.1110: f32[32,4,128,4096], param_1.1273: f32[], param_2.1104: f32[], param_3.793: f32[], param_4.494: f32[32,4,128,4096], param_5.419: f32[], param_6.291: f32[4,32,128,4096], param_7.190: pred[], param_8.108: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { + %param_0.1110 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) + %param_3.793 = f32[]{:T(128)S(6)} parameter(3) + %mul.1519.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.190 = pred[]{:T(512)S(6)} parameter(7) + %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.190), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.291 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.415.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.291), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.419 = f32[]{:T(128)} parameter(5) + %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.419), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.415.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.415.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.931.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.566.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.931.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %mul.1519.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.104 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) + %mul.1525.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.108 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) %constant.935.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1520.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1518.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.104, %mul.1520.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1519.clone.1, %mul.1518.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1095 = f32[]{:T(128)S(6)} parameter(2) - %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1095), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1526.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1524.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.108, %mul.1526.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1525.clone.1, %mul.1524.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1104 = f32[]{:T(128)S(6)} parameter(2) + %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1104), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.64.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %select_n.257.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.934.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1517.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1515.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1517.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.489 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) + %mul.1523.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1521.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1523.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.494 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) %constant.933.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1516.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1514.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.489, %mul.1516.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1515.clone.1, %mul.1514.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1259 = f32[]{:T(128)S(6)} parameter(1) - %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1259), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1522.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1520.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.494, %mul.1522.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1521.clone.1, %mul.1520.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1273 = f32[]{:T(128)S(6)} parameter(1) + %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1273), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.751.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.798.clone.1, %div.752.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.62.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} sqrt(%div.751.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.932.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -638,10 +638,10 @@ StackFrames %add.796.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%sqrt.62.clone.1, %add.797.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.260.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%div.753.clone.1, %add.796.clone.1), metadata={op_name="multiply.38"} %div.750.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.799.clone.1, %multiply.260.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1512.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1096, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1512.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1511.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1513.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1096, %mul.1511.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1518.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1110, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1518.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1517.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1519.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1110, %mul.1517.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.181 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%add.794.clone.1, %add.794.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.997 = f32[]{:T(128)} constant(0) %reduce.136 = f32[]{:T(128)} reduce(%square.181, %constant.997), dimensions={0,1,2,3}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -655,23 +655,23 @@ StackFrames ROOT %reduce_sum.267 = f32[]{:T(128)} add(%reduce_sum.262, %reduce_sum.266), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.280 (param_0.1115: bf16[4,128,128256], param_1.1274: f32[4,128], param_2.1106: s32[4,128], param_3.799: bf16[4,128]) -> f32[4,128] { - %param_2.1106 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1106), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.279 (param_0.1129: bf16[4,128,128256], param_1.1288: f32[4,128], param_2.1115: s32[4,128], param_3.803: bf16[4,128]) -> f32[4,128] { + %param_2.1115 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1115), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.25 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.24 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1115 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.938 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1115), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.799 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.799), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.938, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1274 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_0.1129 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.950 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1129), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_3.803 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) + %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.803), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.950, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.60 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %constant.1017 = f32[]{:T(128)} constant(0) %broadcast.511 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%constant.1017), dimensions={}, metadata={op_name="broadcast.83"} - %mul.1367 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1367, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1373 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1373, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_7.10 (reduce_sum.93: f32[], reduce_sum.94: f32[]) -> f32[] { @@ -680,12 +680,12 @@ StackFrames ROOT %reduce_sum.95 = f32[]{:T(128)} add(%reduce_sum.93, %reduce_sum.94), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.285 (param_0.1116: bf16[4,128,128256], param_1.1275: bf16[4,128]) -> f32[4,128] { - %param_0.1116 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.944 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1116), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1275 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1275), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.944, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.284 (param_0.1130: bf16[4,128,128256], param_1.1289: bf16[4,128]) -> f32[4,128] { + %param_0.1130 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.956 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1130), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1289 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1289), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.956, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.54 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} %constant.1018 = f32[]{:T(128)} constant(0) ROOT %reduce.138 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1018), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} @@ -703,23 +703,23 @@ StackFrames ROOT %reduce_sum.171 = f32[]{:T(128)} add(%reduce_sum.169, %reduce_sum.170), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.291 (param_0.1108: f32[4,4096,8,128], param_1.1268: f32[4,4096,8,128]) -> (f32[], f32[]) { - %param_0.1108 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.344 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.184 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.344, %bitcast.344), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.290 (param_0.1122: f32[4,4096,8,128], param_1.1282: f32[4,4096,8,128]) -> (f32[], f32[]) { + %param_0.1122 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.350 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1122), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.184 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.350, %bitcast.350), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1009 = f32[]{:T(128)} constant(0) %reduce.141 = f32[]{:T(128)} reduce(%square.184, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1268 = f32[4,4096,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) - %bitcast.348.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.187.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.348.clone.1, %bitcast.348.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %param_1.1282 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(1) + %bitcast.354.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.187.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.354.clone.1, %bitcast.354.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %reduce.142.clone.1 = f32[]{:T(128)} reduce(%square.187.clone.1, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.156 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.141, %reduce.142.clone.1) } -%fused_computation.294 (param_0.794: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.794 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.794), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.349 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.293 (param_0.807: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { + %param_0.807 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.807), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} + ROOT %bitcast.355 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_58.63 (reduce_sum.324: f32[], reduce_sum.325: f32[]) -> f32[] { @@ -734,39 +734,39 @@ StackFrames ROOT %reduce_sum.252 = f32[]{:T(128)} add(%reduce_sum.247, %reduce_sum.248), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.295 (param_0.1094: f32[4096,4,8,128], param_1.1257: f32[], param_2.1093: f32[], param_3.787: f32[], param_4.487: f32[4096,4,8,128], param_5.409: f32[], param_6.277: f32[4,4096,8,128], param_7.176: pred[], param_8.102: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1094 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %param_3.787 = f32[]{:T(128)S(6)} parameter(3) - %mul.1496.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.787), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.176 = pred[]{:T(512)S(6)} parameter(7) - %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.176), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.277 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.405.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.277), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.409 = f32[]{:T(128)} parameter(5) - %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.405.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.405.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.294 (param_0.1108: f32[4096,4,8,128], param_1.1271: f32[], param_2.1102: f32[], param_3.791: f32[], param_4.492: f32[4096,4,8,128], param_5.417: f32[], param_6.289: f32[4,4096,8,128], param_7.188: pred[], param_8.106: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { + %param_0.1108 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.791 = f32[]{:T(128)S(6)} parameter(3) + %mul.1502.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.188 = pred[]{:T(512)S(6)} parameter(7) + %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.188), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.289 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.289), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.417 = f32[]{:T(128)} parameter(5) + %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.411.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.919.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.562.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.919.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1500.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1506.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.106 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.923.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.561.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.923.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1499.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.102, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1500.clone.1, %mul.1499.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1093 = f32[]{:T(128)S(6)} parameter(2) - %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1093), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1505.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.106, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1506.clone.1, %mul.1505.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) + %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.62.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %select_n.249.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.922.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.560.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.922.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1498.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.487 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1504.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.492 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.921.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.559.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.921.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1497.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.487, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1498.clone.1, %mul.1497.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1257 = f32[]{:T(128)S(6)} parameter(1) - %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1257), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1503.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.492, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1504.clone.1, %mul.1503.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1271 = f32[]{:T(128)S(6)} parameter(1) + %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1271), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.735.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.786.clone.1, %div.736.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.60.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.735.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.920.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -774,15 +774,15 @@ StackFrames %add.785.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.60.clone.1, %broadcast.557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.258.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.737.clone.1, %add.785.clone.1), metadata={op_name="multiply.40"} %div.734.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.787.clone.1, %multiply.258.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1495.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1094, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1495.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1494.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1496.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1094, %mul.1494.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1501.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1108, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1501.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1500.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1502.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1108, %mul.1500.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.188 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.783.clone.1, %add.783.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.995 = f32[]{:T(128)} constant(0) %reduce.143 = f32[]{:T(128)} reduce(%square.188, %constant.995), dimensions={0,1,2,3}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} %reduce.145.clone.1 = f32[]{:T(128)} reduce(%integer_pow.62.clone.1, %constant.995), dimensions={0,1,2,3}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) + ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) } %region_55.60 (reduce_sum.304: f32[], reduce_sum.308: f32[]) -> f32[] { @@ -797,39 +797,39 @@ StackFrames ROOT %reduce_sum.234 = f32[]{:T(128)} add(%reduce_sum.232, %reduce_sum.233), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.296 (param_0.1097: f32[4096,4,8,128], param_1.1260: f32[], param_2.1096: f32[], param_3.790: f32[], param_4.490: f32[4096,4,8,128], param_5.412: f32[], param_6.280: f32[4,4096,8,128], param_7.179: pred[], param_8.105: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1097 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.790 = f32[]{:T(128)S(6)} parameter(3) - %mul.1523.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.179 = pred[]{:T(512)S(6)} parameter(7) - %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.179), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.280 = f32[4,4096,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.412 = f32[]{:T(128)} parameter(5) - %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.411.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.295 (param_0.1111: f32[4096,4,8,128], param_1.1274: f32[], param_2.1105: f32[], param_3.794: f32[], param_4.495: f32[4096,4,8,128], param_5.420: f32[], param_6.292: f32[4,4096,8,128], param_7.191: pred[], param_8.109: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { + %param_0.1111 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.794 = f32[]{:T(128)S(6)} parameter(3) + %mul.1529.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.191 = pred[]{:T(512)S(6)} parameter(7) + %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.191), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.292 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.417.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.292), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.420 = f32[]{:T(128)} parameter(5) + %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.420), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.417.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.417.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.937.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.572.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.937.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1527.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.105 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %mul.1533.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.109 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.941.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.571.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.941.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1526.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.105, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1527.clone.1, %mul.1526.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1096 = f32[]{:T(128)S(6)} parameter(2) - %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1096), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1532.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.109, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1533.clone.1, %mul.1532.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1105 = f32[]{:T(128)S(6)} parameter(2) + %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1105), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.65.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %select_n.261.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.940.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.570.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.940.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1525.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.490 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %mul.1531.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.495 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.939.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.569.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.939.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1524.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.490, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1525.clone.1, %mul.1524.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1260 = f32[]{:T(128)S(6)} parameter(1) - %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1260), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1530.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.495, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1531.clone.1, %mul.1530.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1274 = f32[]{:T(128)S(6)} parameter(1) + %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.759.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.803.clone.1, %div.760.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.63.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.759.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.938.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -837,10 +837,10 @@ StackFrames %add.802.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.761.clone.1, %add.802.clone.1), metadata={op_name="multiply.37"} %div.758.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.804.clone.1, %multiply.261.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1522.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1097, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1522.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1521.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1523.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1097, %mul.1521.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1528.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1111, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1528.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1527.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1529.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1111, %mul.1527.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.189 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.800.clone.1, %add.800.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.998 = f32[]{:T(128)} constant(0) %reduce.144 = f32[]{:T(128)} reduce(%square.189, %constant.998), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -848,22 +848,22 @@ StackFrames ROOT %tuple.143 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.144, %add.800.clone.1, %add.803.clone.1, %add.804.clone.1, %reduce.146.clone.1) } -%fused_computation.312 (param_0.859: bf16[4,128,4096], param_1.928: f32[4,128], param_2.717: f32[4,128], param_3.448: bf16[4,128,4096], param_4.266: bf16[4096]) -> bf16[4,128,4096] { - %param_3.448 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.266 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.371 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.266), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.361 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.448, %dot_general.371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.961 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.717 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1417 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.717), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1409 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.961, %mul.1417), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.859 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.972 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.859), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.928 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1416 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.928), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1415 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.972, %mul.1416), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1409, %mul.1415), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} - ROOT %convert_element_type.959 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.311 (param_0.872: bf16[4,128,4096], param_1.941: f32[4,128], param_2.726: f32[4,128], param_3.452: bf16[4,128,4096], param_4.271: bf16[4096]) -> bf16[4,128,4096] { + %param_3.452 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.271 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %dot_general.375 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.271), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.365 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.452, %dot_general.375), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.973 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_2.726 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %mul.1423 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.726), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1415 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.973, %mul.1423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_0.872 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.984 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.872), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.941 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1422 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.941), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1421 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.984, %mul.1422), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1415, %mul.1421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} + ROOT %convert_element_type.971 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} } %region_5.8 (reduce_sum.87: f32[], reduce_sum.88: f32[]) -> f32[] { @@ -872,10 +872,10 @@ StackFrames ROOT %reduce_sum.92 = f32[]{:T(128)} add(%reduce_sum.87, %reduce_sum.88), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.313 (param_0.1117: bf16[4,128,4096]) -> f32[4,128] { - %param_0.1117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.963 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1117), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.192 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.963, %convert_element_type.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} +%fused_computation.312 (param_0.1131: bf16[4,128,4096]) -> f32[4,128] { + %param_0.1131 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.975 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1131), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.192 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.975, %convert_element_type.975), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} %constant.1019 = f32[]{:T(128)} constant(0) ROOT %reduce.147 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.192, %constant.1019), dimensions={2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } @@ -886,17 +886,17 @@ StackFrames ROOT %reduce_sum.107 = f32[]{:T(128)} add(%reduce_sum.102, %reduce_sum.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.315 (param_0.1112: bf16[4,128,4096], param_1.1271: bf16[4,128,4096], param_2.1104: bf16[4096]) -> f32[4,128] { - %param_0.1112 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.970 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1112), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1104 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.370 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1104), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.360 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1271, %dot_general.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.969 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1413 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.970, %convert_element_type.969), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} +%fused_computation.314 (param_0.1126: bf16[4,128,4096], param_1.1285: bf16[4,128,4096], param_2.1113: bf16[4096]) -> f32[4,128] { + %param_0.1126 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1126), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1285 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.374 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1113), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.364 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1285, %dot_general.374), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.981 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %mul.1419 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %convert_element_type.981), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.1013 = f32[]{:T(128)} constant(0) - ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1413, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1419, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_8.11 (dot_general.182: bf16[], dot_general.183: bf16[]) -> bf16[] { @@ -905,86 +905,86 @@ StackFrames ROOT %add.168 = bf16[]{:T(256)} add(%dot_general.182, %dot_general.183), metadata={op_name="add.54"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.236.clone.clone (param_0.1081: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1081 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1021 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1081), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.443 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1021), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +%fused_computation.235.clone.clone (param_0.1095: f32[4096,128256]) -> bf16[4096,128256,1] { + %param_0.1095 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %convert_element_type.1033 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1095), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + ROOT %bitcast.449 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} } -%fused_computation.281.clone.1.clone.clone (param_0.1082: bf16[4,128,128256], param_1.1247: s32[4,128], param_2.1072: f32[4,128], param_3.778: f32[4,128], param_4.479: bf16[4,128], param_5.401: f32[4,128]) -> bf16[4,128,128256] { - %param_5.401 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1597 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.401), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.778 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1596 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.778), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1082 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1024 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.479 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.479), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1024, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.280.clone.1.clone.clone (param_0.1096: bf16[4,128,128256], param_1.1261: s32[4,128], param_2.1081: f32[4,128], param_3.782: f32[4,128], param_4.484: bf16[4,128], param_5.409: f32[4,128]) -> bf16[4,128,128256] { + %param_5.409 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.1603 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_3.782 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) + %mul.1602 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.782), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1096 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1036 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1096), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.484 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.484), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1036, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.60 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1595 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1596, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1072 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1072), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1595, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1247 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1247), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.1601 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1602, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1081), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1601, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1261 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1261), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.42 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.41 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1023 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1023), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1594 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1597, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1022 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1594), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.316 (param_0.1080: f32[4,128], param_1.1246: bf16[4,128,4096], param_2.1073: f32[4096,128256], param_3.779: bf16[4,128,128256], param_4.480: s32[4,128], param_5.402: f32[4,128], param_6.272: f32[4,128], param_7.171: bf16[4,128], param_8.98: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { - %param_3.779 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.480 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.272 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.171 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.779, %param_4.480, %param_5.402, %param_6.272, %param_7.171, /*index=5*/%param_8.98), kind=kLoop, calls=%fused_computation.281.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1073 = f32[4096,128256]{1,0:T(8,128)} parameter(2) - %fusion.209.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1073), kind=kLoop, calls=%fused_computation.236.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.80.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.209.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %param_1.1246 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1246), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1080 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1428 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1080), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1427 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %mul.1428), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.981 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1427), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.80.clone.1, %convert_element_type.981), metadata={op_name="multiply.206"} + %convert_element_type.1035 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1035), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1600 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1603, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1034 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1600), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.315 (param_0.1094: f32[4,128], param_1.1260: bf16[4,128,4096], param_2.1082: f32[4096,128256], param_3.783: bf16[4,128,128256], param_4.485: s32[4,128], param_5.410: f32[4,128], param_6.284: f32[4,128], param_7.183: bf16[4,128], param_8.102: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { + %param_3.783 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) + %param_4.485 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %param_5.410 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.284 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.183 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) + %param_8.102 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) + %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.783, %param_4.485, %param_5.410, %param_6.284, %param_7.183, /*index=5*/%param_8.102), kind=kLoop, calls=%fused_computation.280.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_2.1082 = f32[4096,128256]{1,0:T(8,128)} parameter(2) + %fusion.219.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1082), kind=kLoop, calls=%fused_computation.235.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %convolution.86.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.219.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} + %param_1.1260 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.994 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1260), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1094 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1434 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1094), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1433 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.994, %mul.1434), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.993 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1433), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.86.clone.1, %convert_element_type.993), metadata={op_name="multiply.206"} %constant.874 = bf16[]{:T(256)} constant(0) %reduce.149 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.252, %constant.874), dimensions={0,1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.80.clone.1) + ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.86.clone.1) } -%fused_computation.324 (param_0.891: f32[64], param_1.961: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.961 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.961), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.891 = f32[64]{0:T(128)S(1)} parameter(0) - %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.891), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} +%fused_computation.323 (param_0.904: f32[64], param_1.974: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.974), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %param_0.904 = f32[64]{0:T(128)S(1)} parameter(0) + %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.904), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %div.618 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.621, %div.619), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.618), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} - %convert_element_type.990 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1002 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %cos.41.clone.1 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.618), metadata={op_name="jit(train_step)/layers/cos" stack_frame_id=0} - %convert_element_type.989.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.990, %convert_element_type.989.clone.1) + %convert_element_type.1001.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1002, %convert_element_type.1001.clone.1) } -%fused_computation.325 (param_0.888: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.888 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.324 (param_0.901: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.901 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.866 = bf16[]{:T(256)} constant(-inf) - %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.888, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.888, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.34 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.38, %pad.37), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } -%fused_computation.326 (param_0.890: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.890 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.325 (param_0.903: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.903 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.865 = bf16[]{:T(256)} constant(-inf) - %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.890, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.890, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.35 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.40, %pad.39), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } @@ -1000,15 +1000,15 @@ StackFrames ROOT %reduce_sum.162 = f32[]{:T(128)} add(%reduce_sum.157, %reduce_sum.161), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.330 (param_0.1109: f32[4,4096], param_1.1269: f32[4,4096]) -> (f32[], f32[]) { - %param_0.1109 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.365 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.195 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.365, %bitcast.365), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.329 (param_0.1123: f32[4,4096], param_1.1283: f32[4,4096]) -> (f32[], f32[]) { + %param_0.1123 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.371 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.195 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1010 = f32[]{:T(128)} constant(0) %reduce.150 = f32[]{:T(128)} reduce(%square.195, %constant.1010), dimensions={0,1}, to_apply=%region_27.32, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1269 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) - %bitcast.369.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.198.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.369.clone.1, %bitcast.369.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %param_1.1283 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) + %bitcast.375.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %square.198.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %reduce.151.clone.1 = f32[]{:T(128)} reduce(%square.198.clone.1, %constant.1010), dimensions={0,1}, to_apply=%region_26.31, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.150, %reduce.151.clone.1) } @@ -1025,39 +1025,39 @@ StackFrames ROOT %reduce_sum.231 = f32[]{:T(128)} add(%reduce_sum.226, %reduce_sum.227), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.333 (param_0.1098: f32[4096,4], param_1.1261: f32[], param_2.1097: f32[], param_3.791: f32[], param_4.491: f32[4096,4], param_5.413: f32[], param_6.281: f32[4,4096], param_7.180: pred[], param_8.106: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1098 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.791 = f32[]{:T(128)S(6)} parameter(3) - %mul.1530.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.180 = pred[]{:T(512)S(6)} parameter(7) - %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.180), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.281 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.413.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.281), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.413 = f32[]{:T(128)} parameter(5) - %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.413), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.413.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.413.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.332 (param_0.1112: f32[4096,4], param_1.1275: f32[], param_2.1106: f32[], param_3.795: f32[], param_4.496: f32[4096,4], param_5.421: f32[], param_6.293: f32[4,4096], param_7.192: pred[], param_8.110: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1112 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.795 = f32[]{:T(128)S(6)} parameter(3) + %mul.1536.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.192 = pred[]{:T(512)S(6)} parameter(7) + %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.192), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.293 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.419.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.293), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.421 = f32[]{:T(128)} parameter(5) + %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.421), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.419.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.419.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.943.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.578.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.943.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1534.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.106 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1540.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.110 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) %constant.947.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.577.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.947.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1533.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.106, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1534.clone.1, %mul.1533.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1097 = f32[]{:T(128)S(6)} parameter(2) - %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1097), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1539.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.110, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1540.clone.1, %mul.1539.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1106 = f32[]{:T(128)S(6)} parameter(2) + %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1106), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.66.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %select_n.265.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.946.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.576.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.946.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1532.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.491 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1538.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.496 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) %constant.945.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.575.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.945.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1531.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.491, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1532.clone.1, %mul.1531.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1261 = f32[]{:T(128)S(6)} parameter(1) - %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1261), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1537.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.496, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1538.clone.1, %mul.1537.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1275 = f32[]{:T(128)S(6)} parameter(1) + %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1275), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.767.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.808.clone.1, %div.768.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.64.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.767.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.944.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1065,10 +1065,10 @@ StackFrames %add.807.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.262.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.769.clone.1, %add.807.clone.1), metadata={op_name="multiply.36"} %div.766.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.809.clone.1, %multiply.262.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1529.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1098, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1529.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1528.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1530.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1098, %mul.1528.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1535.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1112, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1535.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1534.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1536.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1112, %mul.1534.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.199 = f32[4096,4]{0,1:T(4,128)} multiply(%add.805.clone.1, %add.805.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.999 = f32[]{:T(128)} constant(0) %reduce.152 = f32[]{:T(128)} reduce(%square.199, %constant.999), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1088,39 +1088,39 @@ StackFrames ROOT %reduce_sum.225 = f32[]{:T(128)} add(%reduce_sum.220, %reduce_sum.224), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.334 (param_0.1099: f32[4096,4], param_1.1262: f32[], param_2.1098: f32[], param_3.792: f32[], param_4.492: f32[4096,4], param_5.414: f32[], param_6.282: f32[4,4096], param_7.181: pred[], param_8.107: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1099 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.792 = f32[]{:T(128)S(6)} parameter(3) - %mul.1537.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.181 = pred[]{:T(512)S(6)} parameter(7) - %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.181), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.282 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.415.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.414 = f32[]{:T(128)} parameter(5) - %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.414), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.415.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.415.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.333 (param_0.1113: f32[4096,4], param_1.1276: f32[], param_2.1107: f32[], param_3.796: f32[], param_4.497: f32[4096,4], param_5.422: f32[], param_6.294: f32[4,4096], param_7.193: pred[], param_8.111: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1113 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.796 = f32[]{:T(128)S(6)} parameter(3) + %mul.1543.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.193 = pred[]{:T(512)S(6)} parameter(7) + %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.193), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.294 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.421.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.294), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.422 = f32[]{:T(128)} parameter(5) + %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.422), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.421.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.421.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.949.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.584.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.949.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1541.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.107 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %mul.1547.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.111 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) %constant.953.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.583.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.953.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1540.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.107, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1541.clone.1, %mul.1540.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1098 = f32[]{:T(128)S(6)} parameter(2) - %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1098), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1546.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.111, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1547.clone.1, %mul.1546.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1107 = f32[]{:T(128)S(6)} parameter(2) + %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1107), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.67.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %select_n.269.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.952.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.582.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.952.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1539.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.492 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %mul.1545.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.497 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) %constant.951.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.581.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.951.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1538.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.492, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1539.clone.1, %mul.1538.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1262 = f32[]{:T(128)S(6)} parameter(1) - %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1262), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1544.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.497, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1545.clone.1, %mul.1544.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1276 = f32[]{:T(128)S(6)} parameter(1) + %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1276), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.775.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.813.clone.1, %div.776.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.65.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.775.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.950.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1128,10 +1128,10 @@ StackFrames %add.812.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.65.clone.1, %broadcast.579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.263.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.777.clone.1, %add.812.clone.1), metadata={op_name="multiply.35"} %div.774.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.814.clone.1, %multiply.263.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1536.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1099, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1536.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1535.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1537.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1099, %mul.1535.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1542.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1113, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1542.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1541.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1543.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1113, %mul.1541.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.200 = f32[4096,4]{0,1:T(4,128)} multiply(%add.810.clone.1, %add.810.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1000 = f32[]{:T(128)} constant(0) %reduce.153 = f32[]{:T(128)} reduce(%square.200, %constant.1000), dimensions={0,1}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1145,10 +1145,10 @@ StackFrames ROOT %reduce_sum.101 = f32[]{:T(128)} add(%reduce_sum.99, %reduce_sum.100), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.345 (param_0.1113: bf16[4096]) -> f32[] { - %param_0.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.994 = f32[4096]{0:T(1024)} convert(%param_0.1113), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.203 = f32[4096]{0:T(1024)} multiply(%convert_element_type.994, %convert_element_type.994), metadata={op_name="jit(train_step)/square" stack_frame_id=0} +%fused_computation.344 (param_0.1127: bf16[4096]) -> f32[] { + %param_0.1127 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.1006 = f32[4096]{0:T(1024)} convert(%param_0.1127), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %square.203 = f32[4096]{0:T(1024)} multiply(%convert_element_type.1006, %convert_element_type.1006), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1014 = f32[]{:T(128)} constant(0) ROOT %reduce.156 = f32[]{:T(128)} reduce(%square.203, %constant.1014), dimensions={0}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } @@ -1165,39 +1165,39 @@ StackFrames ROOT %reduce_sum.204 = f32[]{:T(128)} add(%reduce_sum.199, %reduce_sum.203), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.346 (param_0.1103: f32[4096], param_1.1266: f32[], param_2.1102: f32[], param_3.796: f32[], param_4.496: f32[4096], param_5.418: f32[], param_6.286: bf16[4096], param_7.185: pred[], param_8.111: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { - %param_0.1103 = f32[4096]{0:T(1024)S(1)} parameter(0) - %param_3.796 = f32[]{:T(128)S(6)} parameter(3) - %mul.1568.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.185 = pred[]{:T(512)S(6)} parameter(7) - %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.185), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.286 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1009.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.286), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.418 = f32[]{:T(128)} parameter(5) - %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1009.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1009.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} +%fused_computation.345 (param_0.1117: f32[4096], param_1.1280: f32[], param_2.1111: f32[], param_3.800: f32[], param_4.501: f32[4096], param_5.426: f32[], param_6.298: bf16[4096], param_7.197: pred[], param_8.115: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { + %param_0.1117 = f32[4096]{0:T(1024)S(1)} parameter(0) + %param_3.800 = f32[]{:T(128)S(6)} parameter(3) + %mul.1574.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.800), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.197 = pred[]{:T(512)S(6)} parameter(7) + %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.298 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1021.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.298), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_5.426 = f32[]{:T(128)} parameter(5) + %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.426), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1021.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1021.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.973.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.600.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.973.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.1574.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.111 = f32[4096]{0:T(1024)S(1)} parameter(8) + %mul.1580.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.115 = f32[4096]{0:T(1024)S(1)} parameter(8) %constant.977.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1575.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1573.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.111, %mul.1575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1574.clone.1, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) - %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1581.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1579.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.115, %mul.1581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1580.clone.1, %mul.1579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1111 = f32[]{:T(128)S(6)} parameter(2) + %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1111), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.71.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %select_n.285.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.976.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1572.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1570.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.496 = f32[4096]{0:T(1024)S(1)} parameter(4) + %mul.1578.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1576.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.501 = f32[4096]{0:T(1024)S(1)} parameter(4) %constant.975.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1571.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1569.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.496, %mul.1571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1570.clone.1, %mul.1569.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1266 = f32[]{:T(128)S(6)} parameter(1) - %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1266), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1577.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1575.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.501, %mul.1577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1576.clone.1, %mul.1575.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1280 = f32[]{:T(128)S(6)} parameter(1) + %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1280), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.807.clone.1 = f32[4096]{0:T(1024)} divide(%add.835.clone.1, %div.808.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.69.clone.1 = f32[4096]{0:T(1024)} sqrt(%div.807.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.974.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1205,10 +1205,10 @@ StackFrames %add.833.clone.1 = f32[4096]{0:T(1024)} add(%sqrt.69.clone.1, %add.834.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.267.clone.1 = f32[4096]{0:T(1024)} multiply(%div.809.clone.1, %add.833.clone.1), metadata={op_name="multiply.31"} %div.806.clone.1 = f32[4096]{0:T(1024)} divide(%add.836.clone.1, %multiply.267.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1567.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1103, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1566.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1568.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1103, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1573.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1117, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1572.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1574.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1117, %mul.1572.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.204 = f32[4096]{0:T(1024)} multiply(%add.831.clone.1, %add.831.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} %constant.1004 = f32[]{:T(128)} constant(0) %reduce.157 = f32[]{:T(128)} reduce(%square.204, %constant.1004), dimensions={0}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} @@ -1216,27 +1216,27 @@ StackFrames ROOT %tuple.148 = (f32[]{:T(128)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.157, %add.831.clone.1, %add.835.clone.1, %add.836.clone.1, %reduce.158.clone.1) } -%fused_computation.352 (param_0.951: s32[512]) -> s32[1024] { +%fused_computation.351 (param_0.964: s32[512]) -> s32[1024] { %constant.801 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.539 = s32[1024]{0:T(1024)} broadcast(%constant.801), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.951 = s32[512]{0:T(512)S(1)} parameter(0) + %param_0.964 = s32[512]{0:T(512)S(1)} parameter(0) %constant.802 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.951, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.964, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %constant.800 = s32[] constant(128255), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.538 = s32[1024]{0:T(1024)} broadcast(%constant.800), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.539, %pad.41, %broadcast.538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } -%fused_computation.353 (param_0.950: s32[4,128]) -> s32[512] { - %param_0.950 = s32[4,128]{1,0:T(4,128)} parameter(0) +%fused_computation.352 (param_0.963: s32[4,128]) -> s32[512] { + %param_0.963 = s32[4,128]{1,0:T(4,128)} parameter(0) %constant.888 = s32[]{:T(128)} constant(0) %broadcast.546 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.888), dimensions={}, metadata={op_name="broadcast.81"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.950, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} + %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.963, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} %constant.875 = s32[]{:T(128)} constant(128256) %add.760 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.875), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.950, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.950), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} - ROOT %bitcast.370 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} + %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.963, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} + %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} + ROOT %bitcast.376 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} } %region_61.66 (reduce_sum.345: f32[], reduce_sum.346: f32[]) -> f32[] { @@ -1251,52 +1251,52 @@ StackFrames ROOT %reduce_sum.273 = f32[]{:T(128)} add(%reduce_sum.268, %reduce_sum.269), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.354 (param_0.1114: bf16[4,128], param_1.1273: f32[4,128], param_2.1105: f32[4,128], param_3.798: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { - %param_3.798 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) +%fused_computation.353 (param_0.1128: bf16[4,128], param_1.1287: f32[4,128], param_2.1114: f32[4,128], param_3.802: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { + %param_3.802 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) %constant.979.clone.1 = s32[]{:T(128)} constant(0) %broadcast.601.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.979.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.798, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1273), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1114 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1114), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.802, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %param_1.1287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1287), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_0.1128 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1128), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} %add.762 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} %square.207 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %add.762), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} %constant.1016 = f32[]{:T(128)} constant(0) %broadcast.543 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1016), dimensions={}, metadata={op_name="broadcast.32"} - %mul.1467 = f32[4,128]{1,0:T(4,128)} multiply(%square.207, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %mul.1459 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1467, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.159 = f32[]{:T(128)} reduce(%mul.1459, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1105 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1105), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} - %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1467), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %mul.1460.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1460.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %mul.1465.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1473 = f32[4,128]{1,0:T(4,128)} multiply(%square.207, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1465 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1473, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.159 = f32[]{:T(128)} reduce(%mul.1465, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_2.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1114), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1473), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1466.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1466.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %mul.1471.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %constant.891.clone.1 = f32[]{:T(128)} constant(1) %add.757.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.891.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1465.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1471.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} ROOT %tuple.149 = (f32[]{:T(128)}, f32[]{:T(128)}, pred[4,128]{1,0:T(4,128)(4,1)S(1)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%reduce.159, %reduce.160.clone.1, %ne.6.clone.1, %add.750.clone.1) } -%fused_computation.357 (param_0.974: f32[4,128], param_1.1088: f32[4,128]) -> f32[4,128] { - %param_0.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1088 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.356 (param_0.987: f32[4,128], param_1.1101: f32[4,128]) -> f32[4,128] { + %param_0.987 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1101 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.869 = f32[]{:T(128)} constant(0.000244140625) %broadcast.549 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.869), dimensions={}, metadata={op_name="broadcast.264"} - %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1088, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1101, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.867 = f32[]{:T(128)} constant(1e-05) %add.770 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.867), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.769 = f32[4,128]{1,0:T(4,128)} add(%div.656, %add.770), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %rsqrt.90 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} %div.649 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.90, %add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.864 = f32[]{:T(128)} constant(-0.5) - %mul.1471 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1464 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1471), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1463 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.974, %mul.1464), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1477 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1470 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1477), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1469 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.987, %mul.1470), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.863 = f32[]{:T(128)} constant(0.00048828125) - %mul.1470 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - ROOT %mul.1462 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1463, %mul.1470), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1476 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + ROOT %mul.1468 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1469, %mul.1476), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} } %region_0.1 (reduce_sum.67: s32[], reduce_sum.71: s32[]) -> s32[] { @@ -1305,64 +1305,64 @@ StackFrames ROOT %reduce_sum.72 = s32[]{:T(128)} add(%reduce_sum.67, %reduce_sum.71), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.361 (param_0.991: pred[4,128]) -> s32[] { - %param_0.991 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1001 = s32[4,128]{1,0:T(4,128)} convert(%param_0.991), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} +%fused_computation.360 (param_0.1004: pred[4,128]) -> s32[] { + %param_0.1004 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %convert_element_type.1013 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1004), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} %constant.889 = s32[]{:T(128)} constant(0) - ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1001, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1013, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.362 (param_0.976: f32[4,128]) -> f32[4,128] { - %param_0.976 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.361 (param_0.989: f32[4,128]) -> f32[4,128] { + %param_0.989 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.870 = f32[]{:T(128)} constant(0.000244140625) %broadcast.541 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.870), dimensions={}, metadata={op_name="broadcast.264"} - %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.976, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.989, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.868 = f32[]{:T(128)} constant(1e-05) %add.759 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.868), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.756 = f32[4,128]{1,0:T(4,128)} add(%div.654, %add.759), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} ROOT %rsqrt.88 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.756), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} } -%fused_computation.363 (param_0.977: pred[4,128], param_1.1272: f32[]) -> f32[4,128] { - %param_0.977 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} +%fused_computation.362 (param_0.990: pred[4,128], param_1.1286: f32[]) -> f32[4,128] { + %param_0.990 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %param_1.1286 = f32[]{:T(128)S(6)} parameter(1) + %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1286), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} %constant.1015 = f32[]{:T(128)} constant(0) %broadcast.545 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1015), dimensions={}, metadata={op_name="broadcast.32"} - ROOT %mul.1472 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.977, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %mul.1478 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.990, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} } -%fused_computation.365 () -> f32[64] { +%fused_computation.364 () -> f32[64] { %constant.873 = f32[]{:T(128)} constant(500000) %broadcast.552 = f32[64]{0:T(128)} broadcast(%constant.873), dimensions={}, metadata={op_name="broadcast.255"} %iota.46 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/layers/iota" stack_frame_id=0} %constant.872 = s32[]{:T(128)} constant(2) %broadcast.551 = s32[64]{0:T(128)} broadcast(%constant.872), dimensions={}, metadata={op_name="broadcast.256"} - %mul.1473 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} - %convert_element_type.1002 = f32[64]{0:T(128)} convert(%mul.1473), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %mul.1479 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} + %convert_element_type.1014 = f32[64]{0:T(128)} convert(%mul.1479), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %constant.871 = f32[]{:T(128)} constant(0.0078125) %broadcast.550 = f32[64]{0:T(128)} broadcast(%constant.871), dimensions={}, metadata={op_name="broadcast.257"} - %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1002, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1014, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.552, %div.657), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} } -%fused_computation.366 (param_0.989: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.989 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1003 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.989), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %bitcast.371 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.371, %convert_element_type.1003) +%fused_computation.365 (param_0.1002: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { + %param_0.1002 = s32[4,128]{1,0:T(4,128)} parameter(0) + %convert_element_type.1015 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1002), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} + %bitcast.377 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1015), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.377, %convert_element_type.1015) } -%fused_computation.369 (param_0.1089: f32[4096,4]) -> bf16[4,4096] { - %param_0.1089 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.445 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1089), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.445) +%fused_computation.369 (param_0.1103: f32[4096,4]) -> bf16[4,4096] { + %param_0.1103 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.451 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1103), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.451) } -%fused_computation.370 (param_0.1090: f32[4096,4]) -> bf16[4,4096] { - %param_0.1090 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.446 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.446) +%fused_computation.370 (param_0.1104: f32[4096,4]) -> bf16[4,4096] { + %param_0.1104 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.452 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1104), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.452) } %region_6.9 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { @@ -1371,41 +1371,41 @@ StackFrames ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.238.clone.clone (param_0.1076: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1076 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1014 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1076), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.441 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1014), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} -} - -%fused_computation.318.clone.clone (param_0.1077: f32[4,128], param_1.1243: bf16[4,128,4096], param_2.1068: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1068 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.379 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1068), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1243 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1016 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1243), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1077 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1589 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1077), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1588 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1016, %mul.1589), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1015 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1588), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.378 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.379, %convert_element_type.1015), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.371 (param_0.1091: f32[4096,128256], param_1.1254: f32[4,128], param_2.1090: bf16[4,128,4096], param_3.784: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { - %param_1.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1090 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.784 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.230.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1254, %param_2.1090, %param_3.784), kind=kLoop, calls=%fused_computation.318.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1091 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %fusion.211.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1091), kind=kLoop, calls=%fused_computation.238.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.81.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.230.clone.1, %fusion.211.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} +%fused_computation.237.clone.clone (param_0.1090: f32[4096,128256]) -> bf16[4096,128256,1] { + %param_0.1090 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %convert_element_type.1026 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + ROOT %bitcast.447 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1026), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +} + +%fused_computation.317.clone.clone (param_0.1091: f32[4,128], param_1.1257: bf16[4,128,4096], param_2.1077: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1077 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1077), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1257 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1028 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1257), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1091 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1595 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1091), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1594 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1028, %mul.1595), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1027 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1594), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1027), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.371 (param_0.1105: f32[4096,128256], param_1.1268: f32[4,128], param_2.1099: bf16[4,128,4096], param_3.788: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { + %param_1.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1099 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.788 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.240.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1268, %param_2.1099, %param_3.788), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1105 = f32[4096,128256]{1,0:T(8,128)} parameter(0) + %fusion.221.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1105), kind=kLoop, calls=%fused_computation.237.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} + %convolution.87.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.240.clone.1, %fusion.221.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} %constant.992 = bf16[]{:T(256)} constant(-inf) - %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.81.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} - ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.81.clone.1) + %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.87.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.87.clone.1) } -%fused_computation.372 (param_0.1088: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.1088 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %bitcast.444 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1088), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.444) +%fused_computation.372 (param_0.1102: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { + %param_0.1102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %bitcast.450 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1102), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.450) } %convert_element_type.525.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { @@ -1414,13 +1414,13 @@ StackFrames ROOT %add.624 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.121.clone.clone (param_0.1229: bf16[4,4096], param_1.1363: s32[]) -> bf16[4096] { - %param_0.1229 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1363 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.121.clone.clone (param_0.1242: bf16[4,4096], param_1.1376: s32[]) -> bf16[4096] { + %param_0.1242 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) %constant.1116 = s32[]{:T(128)} constant(0) - %dynamic_slice.310 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1229, %param_1.1363, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %dynamic_slice.316 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1242, %param_1.1376, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} %constant.1117 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.310, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.316, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %region_12.14 (reduce_sum.108: f32[], reduce_sum.109: f32[]) -> f32[] { @@ -1429,70 +1429,70 @@ StackFrames ROOT %reduce_sum.113 = f32[]{:T(128)} add(%reduce_sum.108, %reduce_sum.109), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.58.clone.clone (param_0.1230: bf16[4,4,128,4096], param_1.1364: s32[]) -> f32[4,128] { - %param_0.1230 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1364 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.58.clone.clone (param_0.1243: bf16[4,4,128,4096], param_1.1377: s32[]) -> f32[4,128] { + %param_0.1243 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1377 = s32[]{:T(128)S(6)} parameter(1) %constant.1118 = s32[]{:T(128)} constant(0) - %dynamic_slice.311 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1230, %param_1.1364, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.543 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.311), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1081 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.543), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.214 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1081, %convert_element_type.1081), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %dynamic_slice.317 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1243, %param_1.1377, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.548 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1093 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.548), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.214 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1093, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} %constant.1119 = f32[]{:T(128)} constant(0) ROOT %reduce.175 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.214, %constant.1119), dimensions={2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} } -%fused_computation.143.clone.1.clone (param_0.1231: f32[4,128]) -> f32[4,128] { - %param_0.1231 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.143.clone.1.clone (param_0.1244: f32[4,128]) -> f32[4,128] { + %param_0.1244 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1121 = f32[]{:T(128)} constant(0.000244140625) %closed_call.81 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1121), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1231, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1244, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1120 = f32[]{:T(128)} constant(1e-05) %closed_call.80 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1120), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.858 = f32[4,128]{1,0:T(4,128)} add(%div.842, %closed_call.80), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.97 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.858), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.24.clone.1.clone.clone (param_0.1245: bf16[4,4096,32,128], param_1.1374: s32[]) -> bf16[4096,32,128,1] { - %param_0.1245 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1374 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.24.clone.1.clone.clone (param_0.1258: bf16[4,4096,32,128], param_1.1387: s32[]) -> bf16[4096,32,128,1] { + %param_0.1258 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) %constant.1134 = s32[]{:T(128)} constant(0) - %dynamic_slice.317 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1245, %param_1.1374, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.554 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.323 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1258, %param_1.1387, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.559 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.91.clone.clone (param_0.1246: f32[4,128], param_1.1375: bf16[4,4,128,4096], param_2.1167: s32[], param_3.843: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.843 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.843), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1375 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1167 = s32[]{:T(128)S(6)} parameter(2) +%fused_computation.91.clone.clone (param_0.1259: f32[4,128], param_1.1388: bf16[4,4,128,4096], param_2.1176: s32[], param_3.847: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.847 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.847), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1388 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1176 = s32[]{:T(128)S(6)} parameter(2) %constant.1135 = s32[]{:T(128)} constant(0) - %dynamic_slice.318 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1375, %param_2.1167, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.556 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.318), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1089 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.556), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1246 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1703 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1246), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1702 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1089, %mul.1703), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1088 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1088), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.555 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.36.clone.clone (param_0.1247: bf16[4,4096,32,128], param_1.1376: s32[], param_2.1168: f32[4,128], param_3.844: bf16[4,4,128,4096], param_4.525: bf16[4096]) -> bf16[4,128,32,128] { - %param_2.1168 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.844 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) - %param_4.525 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.332 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1168, %param_3.844, %param_1.1376, %param_4.525), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1247 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.331 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1247, %param_1.1376), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.107 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.332, %fusion.331), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.70.clone.clone (param_0.1248: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { - %param_0.1248 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1248), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %dynamic_slice.324 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1388, %param_2.1176, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.561 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1101 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1259 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1709 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1259), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1708 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1101, %mul.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1100 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1708), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1100), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.560 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.427), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.36.clone.clone (param_0.1260: bf16[4,4096,32,128], param_1.1389: s32[], param_2.1177: f32[4,128], param_3.848: bf16[4,4,128,4096], param_4.530: bf16[4096]) -> bf16[4,128,32,128] { + %param_2.1177 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.848 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1389 = s32[]{:T(128)S(6)} parameter(1) + %param_4.530 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.343 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1177, %param_3.848, %param_1.1389, %param_4.530), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1260 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.342 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1260, %param_1.1389), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.113 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.343, %fusion.342), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.70.clone.clone (param_0.1261: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { + %param_0.1261 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.129 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1248), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.187 = (bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) } @@ -1502,172 +1502,172 @@ StackFrames %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} %constant.1123 = s32[]{:T(128)} constant(2) %closed_call.83 = s32[64]{0:T(128)} broadcast(%constant.1123), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1693 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1082 = f32[64]{0:T(128)} convert(%mul.1693), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %mul.1699 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1094 = f32[64]{0:T(128)} convert(%mul.1699), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %constant.1122 = f32[]{:T(128)} constant(0.0078125) %closed_call.82 = f32[64]{0:T(128)} broadcast(%constant.1122), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1082, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1094, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.84, %div.843), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} } -%fused_computation.117.clone.clone (param_0.1232: f32[64], param_1.1365: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1365 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1365), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1232 = f32[64]{0:T(128)S(1)} parameter(0) - %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1232), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} +%fused_computation.117.clone.clone (param_0.1245: f32[64], param_1.1378: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1378 = f32[4,128]{1,0:T(4,128)} parameter(1) + %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1378), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %param_0.1245 = f32[64]{0:T(128)S(1)} parameter(0) + %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1245), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %div.844 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.846, %div.845), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1083 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1095 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} %convert_element_type.829.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1083, %convert_element_type.829.clone.3) + ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1095, %convert_element_type.829.clone.3) } -%fused_computation.120.clone.clone (param_0.1239: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1239 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.120.clone.clone (param_0.1252: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1252 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1130 = bf16[]{:T(256)} constant(-inf) - %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1239, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1239, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.61, %pad.60), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.554 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.119.clone.clone (param_0.1233: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1233 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.119.clone.clone (param_0.1246: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1246 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1125 = bf16[]{:T(256)} constant(-inf) - %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1233, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1233, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.44 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.59, %pad.58), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.544 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.73.clone.clone (param_0.1249: bf16[4,128,32,64], param_1.1377: bf16[4,128,32,64], param_2.1169: bf16[4,128,32,128], param_3.845: bf16[4,128,128], param_4.526: bf16[4,128,128]) -> bf16[4,32,128,128] { - %param_2.1169 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.526 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1707 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.526), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1705 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1169, %mul.1707), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1377 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.73.clone.clone (param_0.1262: bf16[4,128,32,64], param_1.1390: bf16[4,128,32,64], param_2.1178: bf16[4,128,32,128], param_3.849: bf16[4,128,128], param_4.531: bf16[4,128,128]) -> bf16[4,32,128,128] { + %param_2.1178 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.531 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1713 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.531), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1711 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1178, %mul.1713), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1390 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1136 = bf16[]{:T(256)} constant(-inf) - %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1377, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1249 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1249, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1390, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1262 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1262, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.47 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.65, %pad.64), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.845 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1706 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.845), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1704 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1706), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1705, %mul.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.557 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.90.clone.clone (param_0.1241: f32[4,128], param_1.1371: bf16[4,4,128,4096], param_2.1164: s32[], param_3.840: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.840 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.422 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.840), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1371 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1164 = s32[]{:T(128)S(6)} parameter(2) + %param_3.849 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1712 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.849), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1710 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1712), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1711, %mul.1710), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.562 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.90.clone.clone (param_0.1254: f32[4,128], param_1.1384: bf16[4,4,128,4096], param_2.1173: s32[], param_3.844: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.844 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.844), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1384 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1173 = s32[]{:T(128)S(6)} parameter(2) %constant.1132 = s32[]{:T(128)} constant(0) - %dynamic_slice.316 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1371, %param_2.1164, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.316), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1087 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1241 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1697 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1241), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1696 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1087, %mul.1697), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1086 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1696), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.421 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.422, %convert_element_type.1086), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.64.clone.1.clone.clone (param_0.1240: bf16[4,4096,8,128], param_1.1370: s32[]) -> bf16[4096,8,128,1] { - %param_0.1240 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1370 = s32[]{:T(128)S(6)} parameter(1) + %dynamic_slice.322 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1384, %param_2.1173, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.557 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1099 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1703 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1254), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1702 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1099, %mul.1703), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1098 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1098), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.556 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.425), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.64.clone.1.clone.clone (param_0.1253: bf16[4,4096,8,128], param_1.1383: s32[]) -> bf16[4096,8,128,1] { + %param_0.1253 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1383 = s32[]{:T(128)S(6)} parameter(1) %constant.1131 = s32[]{:T(128)} constant(0) - %dynamic_slice.315 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1240, %param_1.1370, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.550 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.315), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.321 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1383, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.555 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.89.clone.clone (param_0.1242: bf16[4,4096,8,128], param_1.1372: s32[], param_2.1165: f32[4,128], param_3.841: bf16[4,4,128,4096], param_4.523: bf16[4096]) -> bf16[4,128,8,128] { - %param_2.1165 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.841 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1372 = s32[]{:T(128)S(6)} parameter(1) - %param_4.523 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.329 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1165, %param_3.841, %param_1.1372, %param_4.523), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1242 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.330 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1242, %param_1.1372), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.329, %fusion.330), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.89.clone.clone (param_0.1255: bf16[4,4096,8,128], param_1.1385: s32[], param_2.1174: f32[4,128], param_3.845: bf16[4,4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,8,128] { + %param_2.1174 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.845 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) + %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.340 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1174, %param_3.845, %param_1.1385, %param_4.528), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1255 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.341 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1255, %param_1.1385), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.112 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.340, %fusion.341), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } -%fused_computation.106.clone.clone (param_0.1243: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1243 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1243), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.106.clone.clone (param_0.1256: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { + %param_0.1256 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1243), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.186 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) } -%fused_computation.109.clone.clone (param_0.1244: bf16[4,128,8,64], param_1.1373: bf16[4,128,8,64], param_2.1166: bf16[4,128,8,128], param_3.842: bf16[4,128,128], param_4.524: bf16[4,128,128]) -> bf16[4,8,128,128] { - %param_2.1166 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.524 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1701 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.524), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1699 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1166, %mul.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1373 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.109.clone.clone (param_0.1257: bf16[4,128,8,64], param_1.1386: bf16[4,128,8,64], param_2.1175: bf16[4,128,8,128], param_3.846: bf16[4,128,128], param_4.529: bf16[4,128,128]) -> bf16[4,8,128,128] { + %param_2.1175 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.529 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1707 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.529), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1705 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1175, %mul.1707), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1386 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1133 = bf16[]{:T(256)} constant(-inf) - %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1373, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1244 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1244, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1386, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1257 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1257, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.46 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.63, %pad.62), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.842 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1700 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1698 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1700), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1699, %mul.1698), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_3.846 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1706 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.846), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1704 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1706), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1705, %mul.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.558 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.135.clone.clone (param_0.1235: bf16[4,4096,8,128], param_1.1367: s32[]) -> bf16[1,4096,8,128] { - %param_0.1235 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1367 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.135.clone.clone (param_0.1248: bf16[4,4096,8,128], param_1.1380: s32[]) -> bf16[1,4096,8,128] { + %param_0.1248 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) + %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) %constant.1128 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.313 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1235, %param_1.1367, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %dynamic_slice.319 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1248, %param_1.1380, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } -%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1236: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { - %param_0.1236 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1236), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.545 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1249: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { + %param_0.1249 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} + ROOT %bitcast.550 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.88.clone.clone.clone.clone (param_0.1237: f32[4,128], param_1.1368: bf16[4,4,128,4096], param_2.1162: s32[], param_3.838: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.838 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.420 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.838), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1368 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1162 = s32[]{:T(128)S(6)} parameter(2) +%fused_computation.88.clone.clone.clone.clone (param_0.1250: f32[4,128], param_1.1381: bf16[4,4,128,4096], param_2.1171: s32[], param_3.842: bf16[4096]) -> bf16[4,128,4096,1] { + %param_3.842 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1381 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) %constant.1129 = s32[]{:T(128)} constant(0) - %dynamic_slice.314 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1368, %param_2.1162, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.547 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.314), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1085 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.547), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1237 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1695 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1237), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1694 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1085, %mul.1695), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1084 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1694), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.419 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.420, %convert_element_type.1084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.546 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.419), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.114.clone.clone (param_0.1238: bf16[1,4096,8,128], param_1.1369: f32[4,128], param_2.1163: bf16[4,4,128,4096], param_3.839: s32[], param_4.522: bf16[4096]) -> bf16[4,8,128,128] { - %param_1.1369 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1163 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.839 = s32[]{:T(128)S(6)} parameter(3) - %param_4.522 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.328 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1369, %param_2.1163, %param_3.839, %param_4.522), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1238 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.327 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1238), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.105 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.328, %fusion.327), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.548 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.146.clone.clone (param_0.1273: f32[4,32,128,128]) -> (f32[4,32,128], f32[4,32,128,1]) { - %param_0.1273 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1273), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.570 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.192 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.570, %slice.11) + %dynamic_slice.320 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1381, %param_2.1171, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1097 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1250 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1701 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1250), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1700 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1097, %mul.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1096 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1700), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1096), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.114.clone.clone (param_0.1251: bf16[1,4096,8,128], param_1.1382: f32[4,128], param_2.1172: bf16[4,4,128,4096], param_3.843: s32[], param_4.527: bf16[4096]) -> bf16[4,8,128,128] { + %param_1.1382 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1172 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.843 = s32[]{:T(128)S(6)} parameter(3) + %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.339 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1382, %param_2.1172, %param_3.843, %param_4.527), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1251 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.338 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1251), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.111 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.339, %fusion.338), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.366.clone.clone (param_0.1286: f32[4,32,128,128]) -> (f32[4,32,128,1], f32[4,32,128]) { + %param_0.1286 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) + %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1286), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.262.clone.3 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} + ROOT %tuple.192 = (f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}, f32[4,32,128]{2,1,0:T(8,128)S(1)}) tuple(%slice.11, %bitcast.262.clone.3) } %region_13.16 (reduce_sum.120: f32[], reduce_sum.121: f32[]) -> f32[] { @@ -1676,34 +1676,34 @@ StackFrames ROOT %reduce_sum.122 = f32[]{:T(128)} add(%reduce_sum.120, %reduce_sum.121), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1250: bf16[4,32,128,4096], param_1.1378: s32[]) -> bf16[32,128,4096,1] { - %param_0.1250 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1378 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1263: bf16[4,32,128,4096], param_1.1391: s32[]) -> bf16[32,128,4096,1] { + %param_0.1263 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) %constant.1137 = s32[]{:T(128)} constant(0) - %dynamic_slice.319 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1250, %param_1.1378, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.558 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.319), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1263, %param_1.1391, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.563 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1251: bf16[4,32,128,128]) -> bf16[4,128,32,128] { - %param_0.1251 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.559 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1264: bf16[4,32,128,128]) -> bf16[4,128,32,128] { + %param_0.1264 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.564 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.61.clone.clone (param_0.1252: bf16[4,32,128,4096], param_1.1379: s32[], param_2.1170: bf16[4,32,128,128], param_3.846: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { - %param_3.846 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.61.clone.clone (param_0.1265: bf16[4,32,128,4096], param_1.1392: s32[], param_2.1179: bf16[4,32,128,128], param_3.850: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { + %param_3.850 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1392 = s32[]{:T(128)S(6)} parameter(1) %constant.365.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.210.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.846, %param_1.1379, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.210.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1170 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.80.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1170), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1252 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.79.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1252, %param_1.1379), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.60.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.80.clone.3, %fusion.79.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %dynamic_slice.208.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.850, %param_1.1392, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.208.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %param_2.1179 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.83.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1179), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_0.1265 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.82.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1265, %param_1.1392), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.62.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.83.clone.3, %fusion.82.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %add.635.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.207.clone.3, %bitcast.182.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1090 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.215 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1090, %convert_element_type.1090), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %convert_element_type.1102 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.215 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1102, %convert_element_type.1102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} %constant.1138 = f32[]{:T(128)} constant(0) %reduce.177 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.215, %constant.1138), dimensions={2}, to_apply=%region_13.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.188 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.177, %add.635.clone.3) @@ -1715,140 +1715,140 @@ StackFrames ROOT %add.623 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.122.clone.clone (param_0.1234: bf16[4,4096], param_1.1366: s32[]) -> bf16[4096] { - %param_0.1234 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1366 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.122.clone.clone (param_0.1247: bf16[4,4096], param_1.1379: s32[]) -> bf16[4096] { + %param_0.1247 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) %constant.1126 = s32[]{:T(128)} constant(0) - %dynamic_slice.312 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1234, %param_1.1366, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %dynamic_slice.318 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1247, %param_1.1379, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} %constant.1127 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.312, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.318, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.12.clone.clone.clone (param_0.1253: bf16[4,14336,4096], param_1.1380: s32[]) -> bf16[14336,4096,1] { - %param_0.1253 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.12.clone.clone.clone (param_0.1266: bf16[4,14336,4096], param_1.1393: s32[]) -> bf16[14336,4096,1] { + %param_0.1266 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1393 = s32[]{:T(128)S(6)} parameter(1) %constant.1139 = s32[]{:T(128)} constant(0) - %dynamic_slice.320 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1380, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.561 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.326 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1266, %param_1.1393, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.566 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.326), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %bitcast_fusion.3.clone.clone (bitcast_input.12: bf16[4,128,4096]) -> bf16[4,128,4096] { %bitcast_input.12 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.560 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) + ROOT %bitcast.565 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) } -%fused_computation.13.clone.clone (param_0.1254: bf16[4,128,4096], param_1.1381: bf16[4,14336,4096], param_2.1171: s32[]) -> bf16[14336,4,128] { - %param_1.1381 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) - %fusion.333 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1381, %param_2.1171), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1254 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %fusion.334 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1254), kind=kLoop, calls=%bitcast_fusion.3.clone.clone - ROOT %convolution.108 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.333, %fusion.334), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.13.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1394: bf16[4,14336,4096], param_2.1180: s32[]) -> bf16[14336,4,128] { + %param_1.1394 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) + %fusion.344 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1394, %param_2.1180), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.345 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1267), kind=kLoop, calls=%bitcast_fusion.3.clone.clone + ROOT %convolution.114 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.344, %fusion.345), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} } -%fused_computation.144.clone.1.clone (param_0.1255: f32[4,128]) -> f32[4,128] { - %param_0.1255 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.144.clone.1.clone (param_0.1268: f32[4,128]) -> f32[4,128] { + %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1141 = f32[]{:T(128)} constant(0.000244140625) %closed_call.86 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1141), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1255, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1140 = f32[]{:T(128)} constant(1e-05) %closed_call.85 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1140), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.861 = f32[4,128]{1,0:T(4,128)} add(%div.847, %closed_call.85), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.98 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.11.clone.1.clone.clone (param_0.1259: bf16[4,4096,14336], param_1.1385: s32[]) -> bf16[4096,14336,1] { - %param_0.1259 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.11.clone.1.clone.clone (param_0.1272: bf16[4,4096,14336], param_1.1398: s32[]) -> bf16[4096,14336,1] { + %param_0.1272 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1398 = s32[]{:T(128)S(6)} parameter(1) %constant.1143 = s32[]{:T(128)} constant(0) - %dynamic_slice.322 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1259, %param_1.1385, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.563 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.328 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1272, %param_1.1398, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.568 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.96.clone.2.clone.clone (param_0.1260: f32[4,128], param_1.1386: bf16[4,128,4096], param_2.1174: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1174), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1094 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1260 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1711 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1260), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1710 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1094, %mul.1711), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1093 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1710), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.96.clone.2.clone.clone (param_0.1273: f32[4,128], param_1.1399: bf16[4,128,4096], param_2.1183: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1183 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.432 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1183), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1399 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1106 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1717 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1273), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1716 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1106, %mul.1717), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1105 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1716), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.431 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.432, %convert_element_type.1105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.23.clone.clone (param_0.1261: bf16[4,4096,14336], param_1.1387: s32[], param_2.1175: f32[4,128], param_3.848: bf16[4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1175 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.848 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.338 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1175, %param_3.848, %param_4.528), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1261 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) - %fusion.337 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1261, %param_1.1387), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.110 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.338, %fusion.337), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.23.clone.clone (param_0.1274: bf16[4,4096,14336], param_1.1400: s32[], param_2.1184: f32[4,128], param_3.852: bf16[4,128,4096], param_4.533: bf16[4096]) -> bf16[4,128,14336] { + %param_2.1184 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.533 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.349 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1184, %param_3.852, %param_4.533), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1274 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1400 = s32[]{:T(128)S(6)} parameter(1) + %fusion.348 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1274, %param_1.1400), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.116 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.349, %fusion.348), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } -%fused_computation.14.clone.1.clone.clone (param_0.1262: bf16[4,4096,14336], param_1.1388: s32[]) -> bf16[4096,14336,1] { - %param_0.1262 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1388 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.14.clone.1.clone.clone (param_0.1275: bf16[4,4096,14336], param_1.1401: s32[]) -> bf16[4096,14336,1] { + %param_0.1275 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1401 = s32[]{:T(128)S(6)} parameter(1) %constant.1144 = s32[]{:T(128)} constant(0) - %dynamic_slice.323 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1262, %param_1.1388, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.564 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.329 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1275, %param_1.1401, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.569 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.329), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.39.clone.1.clone.clone (param_0.1263: bf16[14336,4,128], param_1.1389: bf16[4,128,14336]) -> bf16[4,128,14336] { - %param_1.1389 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) +%fused_computation.39.clone.1.clone.clone (param_0.1276: bf16[14336,4,128], param_1.1402: bf16[4,128,14336]) -> bf16[4,128,14336] { + %param_1.1402 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) %constant.1145 = bf16[]{:T(256)} constant(1) %jit_silu_.44 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1145), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1402), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.69 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.862 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.848 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.862), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1713 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1389, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %param_0.1263 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) - %bitcast.565 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - ROOT %mul.1712 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1713, %bitcast.565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1719 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1402, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %param_0.1276 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.570 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + ROOT %mul.1718 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1719, %bitcast.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} } -%fused_computation.21.clone.clone (param_0.1264: bf16[4,4096,14336], param_1.1390: s32[], param_2.1176: bf16[14336,4,128], param_3.849: bf16[4,128,14336]) -> bf16[4,128,4096] { - %param_2.1176 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %param_3.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1176, %param_3.849), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_0.1264 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1390 = s32[]{:T(128)S(6)} parameter(1) - %fusion.339 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1264, %param_1.1390), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.111 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.339), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.21.clone.clone (param_0.1277: bf16[4,4096,14336], param_1.1403: s32[], param_2.1185: bf16[14336,4,128], param_3.853: bf16[4,128,14336]) -> bf16[4,128,4096] { + %param_2.1185 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) + %param_3.853 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1185, %param_3.853), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %param_0.1277 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1403 = s32[]{:T(128)S(6)} parameter(1) + %fusion.350 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1277, %param_1.1403), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.350), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} } -%fused_computation.14.clone.clone.clone (param_0.1256: bf16[4,4096,14336], param_1.1382: s32[]) -> bf16[4096,14336,1] { - %param_0.1256 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1382 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.14.clone.clone.clone (param_0.1269: bf16[4,4096,14336], param_1.1395: s32[]) -> bf16[4096,14336,1] { + %param_0.1269 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1395 = s32[]{:T(128)S(6)} parameter(1) %constant.1142 = s32[]{:T(128)} constant(0) - %dynamic_slice.321 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1256, %param_1.1382, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.562 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.327 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1269, %param_1.1395, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.567 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.96.clone.1.clone.clone (param_0.1257: f32[4,128], param_1.1383: bf16[4,128,4096], param_2.1172: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1172 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1172), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1092 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1257 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1709 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1257), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1708 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1092, %mul.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1091 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1708), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1091), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.96.clone.1.clone.clone (param_0.1270: f32[4,128], param_1.1396: bf16[4,128,4096], param_2.1181: bf16[4096]) -> bf16[4,128,4096] { + %param_2.1181 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1181), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1396 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1104 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1270 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1715 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1270), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1714 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1104, %mul.1715), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1103 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1714), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.430, %convert_element_type.1103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.20.clone.clone (param_0.1258: bf16[4,4096,14336], param_1.1384: s32[], param_2.1173: f32[4,128], param_3.847: bf16[4,128,4096], param_4.527: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1173 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.847 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.336 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1173, %param_3.847, %param_4.527), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1258 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1384 = s32[]{:T(128)S(6)} parameter(1) - %fusion.335 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1258, %param_1.1384), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.109 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.336, %fusion.335), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +%fused_computation.20.clone.clone (param_0.1271: bf16[4,4096,14336], param_1.1397: s32[], param_2.1182: f32[4,128], param_3.851: bf16[4,128,4096], param_4.532: bf16[4096]) -> bf16[4,128,14336] { + %param_2.1182 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.532 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.347 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1182, %param_3.851, %param_4.532), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1271 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1397 = s32[]{:T(128)S(6)} parameter(1) + %fusion.346 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1271, %param_1.1397), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.115 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.347, %fusion.346), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} } %region_14.17 (reduce_sum.126: f32[], reduce_sum.127: f32[]) -> f32[] { @@ -1857,63 +1857,63 @@ StackFrames ROOT %reduce_sum.128 = f32[]{:T(128)} add(%reduce_sum.126, %reduce_sum.127), metadata={op_name="checkpoint/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1265: bf16[4,4096,14336], param_1.1391: s32[]) -> bf16[4096,14336,1] { - %param_0.1265 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1278: bf16[4,4096,14336], param_1.1404: s32[]) -> bf16[4096,14336,1] { + %param_0.1278 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1404 = s32[]{:T(128)S(6)} parameter(1) %constant.1146 = s32[]{:T(128)} constant(0) - %dynamic_slice.324 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1265, %param_1.1391, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.566 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.330 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1278, %param_1.1404, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.571 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1266: bf16[4,128,14336], param_1.1392: bf16[4,128,14336], param_2.1177: bf16[14336,4,128]) -> bf16[4,128,14336] { - %param_2.1177 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %bitcast.567 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1177), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %param_1.1392 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %mul.1718 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.567, %param_1.1392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} +%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1279: bf16[4,128,14336], param_1.1405: bf16[4,128,14336], param_2.1186: bf16[14336,4,128]) -> bf16[4,128,14336] { + %param_2.1186 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.572 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1186), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + %param_1.1405 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %mul.1724 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.572, %param_1.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1147 = bf16[]{:T(256)} constant(1) %jit_silu_.45 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1147), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %param_0.1266 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1266), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %param_0.1279 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.70 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.131), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.863 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.70, %jit_silu_.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.45, %add.863), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1717 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1718, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1716 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1266, %mul.1718), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1723 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1724, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1722 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1279, %mul.1724), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} %sub.98 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} subtract(%jit_silu_.45, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/sub" stack_frame_id=0} - %mul.1715 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1714 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1716, %mul.1715), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1717, %mul.1714), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} -} - -%fused_computation.63.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1393: bf16[4096], param_2.1178: bf16[4,128,4096], param_3.850: bf16[4,4096,14336], param_4.529: s32[], param_5.425: bf16[4,128,14336], param_6.291: bf16[4,128,14336], param_7.188: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { - %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1096 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1267), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_2.1178 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_5.425 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) - %param_6.291 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) - %param_7.188 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) - %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.425, %param_6.291, %param_7.188), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} - %param_3.850 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.529 = s32[]{:T(128)S(6)} parameter(4) - %fusion.91.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.850, %param_4.529), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.64.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.91.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1178, %convolution.64.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1393 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) - %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1393), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.430), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1095 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.429), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1719 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1096, %convert_element_type.1095), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1721 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %mul.1720 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1722, %mul.1721), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} + ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1723, %mul.1720), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} +} + +%fused_computation.63.clone.clone (param_0.1280: bf16[4,128,4096], param_1.1406: bf16[4096], param_2.1187: bf16[4,128,4096], param_3.854: bf16[4,4096,14336], param_4.534: s32[], param_5.435: bf16[4,128,14336], param_6.304: bf16[4,128,14336], param_7.200: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { + %param_0.1280 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1108 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_2.1187 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_5.435 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) + %param_6.304 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) + %param_7.200 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) + %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.435, %param_6.304, %param_7.200), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} + %param_3.854 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) + %param_4.534 = s32[]{:T(128)S(6)} parameter(4) + %fusion.79.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.854, %param_4.534), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.60.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.79.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} + %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1187, %convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %param_1.1406 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %dot_general.434 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1406), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.433 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.434), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1107 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.433), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} + %mul.1725 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1108, %convert_element_type.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1148 = f32[]{:T(128)} constant(0) - %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1719, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} + %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1725, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.189 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.178, %add_any.132.clone.3) } -%fused_computation.140.clone.clone (param_0.1268: f32[4,128], param_1.1394: f32[4,128]) -> f32[4,128] { - %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1394 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.140.clone.clone (param_0.1281: f32[4,128], param_1.1407: f32[4,128]) -> f32[4,128] { + %param_0.1281 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1407 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.1152 = f32[]{:T(128)} constant(0.000244140625) %closed_call.89 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1152), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1394, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1407, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1151 = f32[]{:T(128)} constant(1e-05) %closed_call.88 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1151), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.864 = f32[4,128]{1,0:T(4,128)} add(%div.851, %closed_call.88), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} @@ -1921,11 +1921,11 @@ StackFrames %div.850 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.99, %add.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %constant.1150 = f32[]{:T(128)} constant(-0.5) %closed_call.87 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1150), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1722 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1721 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %mul.1722), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1728 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1727 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1281, %mul.1728), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %constant.1149 = f32[]{:T(128)} constant(0.00048828125) - %mul.1723 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - ROOT %mul.1720 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1721, %mul.1723), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1729 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + ROOT %mul.1726 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1727, %mul.1729), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} } %region_20.24 (dot_general.187: bf16[], dot_general.188: bf16[]) -> bf16[] { @@ -1934,29 +1934,29 @@ StackFrames ROOT %add.173 = bf16[]{:T(256)} add(%dot_general.187, %dot_general.188), metadata={op_name="add.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.94.clone.clone (param_0.1269: bf16[4,128,4096], param_1.1395: f32[4,128], param_2.1179: bf16[4,128,4096], param_3.851: bf16[4,128,4096], param_4.530: f32[4,128], param_5.426: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { - %param_0.1269 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_2.1179 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %convert_element_type.1098 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1179), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_1.1395 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1725 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1395), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1724 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1098, %mul.1725), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1097 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1724), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1269, %convert_element_type.1097), metadata={op_name="multiply.204"} +%fused_computation.94.clone.clone (param_0.1282: bf16[4,128,4096], param_1.1408: f32[4,128], param_2.1188: bf16[4,128,4096], param_3.855: bf16[4,128,4096], param_4.535: f32[4,128], param_5.436: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { + %param_0.1282 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_2.1188 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1110 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1188), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_1.1408 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1731 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1408), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.1730 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1109 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1730), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %convert_element_type.1109), metadata={op_name="multiply.204"} %constant.1153 = bf16[]{:T(256)} constant(0) %reduce.179 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.271, %constant.1153), dimensions={0,1}, to_apply=%region_20.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_5.426 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) - %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.426), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1269, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.855 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_5.436 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.436), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} %convert_element_type.753.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.263.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1142.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1725), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_4.530 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %mul.1151.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.530), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %mul.1141.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1098, %mul.1151.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1142.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %param_4.535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %mul.1151.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.535), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} + %mul.1141.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1151.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} %add_any.126.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1142.clone.3, %mul.1141.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} %convert_element_type.751.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.126.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.851, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.855, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} ROOT %tuple.190 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.179, %add_any.124.clone.3) } @@ -1966,35 +1966,35 @@ StackFrames ROOT %add.169 = f32[]{:T(128)} add(%dot_general.184, %dot_general.185), metadata={op_name="add.31"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1270: bf16[4,32,128,4096], param_1.1396: s32[]) -> bf16[32,128,4096,1] { - %param_0.1270 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1396 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1283: bf16[4,32,128,4096], param_1.1409: s32[]) -> bf16[32,128,4096,1] { + %param_0.1283 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1409 = s32[]{:T(128)S(6)} parameter(1) %constant.1154 = s32[]{:T(128)} constant(0) - %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1270, %param_1.1396, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.568 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.331 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1283, %param_1.1409, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.573 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1271: bf16[4,128,4096]) -> bf16[4,128,4096,1] { - %param_0.1271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.569 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1271), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} +%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1284: bf16[4,128,4096]) -> bf16[4,128,4096,1] { + %param_0.1284 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.574 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} } -%fused_computation.66.clone.clone (param_0.1272: bf16[4,32,128,128], param_1.1397: bf16[4,32,128,4096], param_2.1180: s32[], param_3.852: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { - %param_0.1272 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1272), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} - %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %fusion.84.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.852), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1397 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) - %fusion.83.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1397, %param_2.1180), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.62.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.84.clone.3, %fusion.83.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +%fused_computation.66.clone.clone (param_0.1285: bf16[4,32,128,128], param_1.1410: bf16[4,32,128,4096], param_2.1189: s32[], param_3.856: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { + %param_0.1285 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} + %param_3.856 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %fusion.95.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.856), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} + %param_1.1410 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1189 = s32[]{:T(128)S(6)} parameter(2) + %fusion.94.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1410, %param_2.1189), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.64.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.95.clone.3, %fusion.94.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} %constant.619.clone.3 = bf16[]{:T(256)} constant(0.25) %div.442.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.619.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} - %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.62.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} + %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.64.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} %bitcast.209.clone.3 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%div.441.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} %convert.123 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.209.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert.1" stack_frame_id=0} %multiply.272 = f32[4,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.124, %convert.123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/multiply" stack_frame_id=0} %constant.1155 = f32[]{:T(128)} constant(0) - %dot_general.431 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} - ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.431, %bitcast.209.clone.3) + %dot_general.435 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} + ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.435, %bitcast.209.clone.3) } diff --git a/tests/utils/reference_hlo_qwen3_1.7b.txt b/tests/utils/reference_hlo_qwen3_1.7b.txt index b54b810621..f1ede66966 100644 --- a/tests/utils/reference_hlo_qwen3_1.7b.txt +++ b/tests/utils/reference_hlo_qwen3_1.7b.txt @@ -32,7 +32,7 @@ StackFrames %param_1.5 = s32[512]{0:T(512)S(1)} parameter(1) %reshape.451 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.5), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} %transpose.466 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.451), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)} parameter(2) + %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)S(1)} parameter(2) %reshape.452 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} reshape(%param_2.4), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} %transpose.467 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} transpose(%reshape.452), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} ROOT %scatter.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.466, %transpose.467), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_42.47.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} @@ -50,43 +50,43 @@ StackFrames ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.386, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.277 (param_0.1367: f32[151936,2048], param_1.1549: f32[], param_2.1311: f32[], param_3.918: f32[], param_4.554: f32[151936,2048], param_5.467: f32[], param_6.356: bf16[151936,2048], param_7.196: bf16[151936,2048,1], param_8.113: pred[], param_9.94: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { - %param_0.1367 = f32[151936,2048]{1,0:T(8,128)} parameter(0) +%fused_computation.277 (param_0.1368: f32[151936,2048], param_1.1556: f32[], param_2.1314: f32[], param_3.918: f32[], param_4.556: f32[151936,2048], param_5.468: f32[], param_6.358: bf16[151936,2048], param_7.201: bf16[151936,2048,1], param_8.118: pred[], param_9.97: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { + %param_0.1368 = f32[151936,2048]{1,0:T(8,128)} parameter(0) %param_3.918 = f32[]{:T(128)S(6)} parameter(3) %mul.1926.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_3.918), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.113 = pred[]{:T(512)S(6)} parameter(8) - %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.113), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_7.196 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) - %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.196), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %param_8.118 = pred[]{:T(512)S(6)} parameter(8) + %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.118), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_7.201 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) + %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.201), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %convert_element_type.1409.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.464.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_6.356 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.356), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_6.358 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %add_any.197.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1409.clone.1, %convert_element_type.1408.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} - %param_5.467 = f32[]{:T(128)} parameter(5) - %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.467), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_5.468 = f32[]{:T(128)} parameter(5) + %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.859.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add_any.197.clone.1, %div.860.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.267.clone.1 = f32[151936,2048]{1,0:T(8,128)} select(%select_n.268.clone.1, %add_any.197.clone.1, %div.859.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1092.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.844.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1092.clone.1), dimensions={}, metadata={op_name="broadcast.74"} %mul.1932.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_9.94 = f32[151936,2048]{1,0:T(8,128)} parameter(9) + %param_9.97 = f32[151936,2048]{1,0:T(8,128)} parameter(9) %constant.1096.clone.1 = f32[]{:T(128)} constant(0.9) %mul.1933.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1096.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1931.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.94, %mul.1933.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1931.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.97, %mul.1933.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.941.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1932.clone.1, %mul.1931.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1311 = f32[]{:T(128)S(6)} parameter(2) - %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1311), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) + %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.65.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %select_n.267.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1095.clone.1 = f32[]{:T(128)} constant(0.05) %mul.1930.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1095.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.1928.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %mul.1930.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.554 = f32[151936,2048]{1,0:T(8,128)} parameter(4) + %param_4.556 = f32[151936,2048]{1,0:T(8,128)} parameter(4) %constant.1094.clone.1 = f32[]{:T(128)} constant(0.95) %mul.1929.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1094.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1927.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.554, %mul.1929.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1927.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.556, %mul.1929.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.940.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1928.clone.1, %mul.1927.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1549 = f32[]{:T(128)S(6)} parameter(1) - %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1549), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) + %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.854.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.940.clone.1, %div.855.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.62.clone.1 = f32[151936,2048]{1,0:T(8,128)} sqrt(%div.854.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1093.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -94,14 +94,14 @@ StackFrames %add.938.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%sqrt.62.clone.1, %add.939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.426.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%div.856.clone.1, %add.938.clone.1), metadata={op_name="multiply.61"} %div.853.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.941.clone.1, %multiply.426.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1925.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1367, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1925.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1368, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.937.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%div.853.clone.1, %mul.1925.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1924.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%mul.1926.clone.1, %add.937.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1367, %mul.1924.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1368, %mul.1924.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.214 = f32[151936,2048]{1,0:T(8,128)} multiply(%add.936.clone.1, %add.936.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1198 = f32[]{:T(128)} constant(0) - %reduce.176 = f32[]{:T(128)} reduce(%square.214, %constant.1198), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1198), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1200 = f32[]{:T(128)} constant(0) + %reduce.176 = f32[]{:T(128)} reduce(%square.214, %constant.1200), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.144 = (f32[]{:T(128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.176, %add.936.clone.1, %add.940.clone.1, %add.941.clone.1, %reduce.178.clone.1) } @@ -111,64 +111,64 @@ StackFrames ROOT %reduce_sum.319 = f32[]{:T(128)} add(%reduce_sum.317, %reduce_sum.318), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.367.clone.clone (param_0.1354: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1287: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1287 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1287), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1445 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1354 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2075 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1354), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2074 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1445, %mul.2075), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1444 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2074), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.478 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.479, %convert_element_type.1444), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.367.clone.clone (param_0.1355: f32[4,128], param_1.1549: bf16[4,128,2048], param_2.1290: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1290 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.480 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1290), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1549 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1451 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1355 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2083 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1355), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.2082 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1451, %mul.2083), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.480, %convert_element_type.1450), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.289.clone.clone.clone (param_0.1355: bf16[4,128,151936], param_1.1543: s32[4,128], param_2.1288: f32[4,128], param_3.911: f32[4,128], param_4.544: bf16[4,128], param_5.445: f32[4,128]) -> bf16[4,128,151936] { - %param_5.445 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2079 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.445), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.289.clone.clone.clone (param_0.1356: bf16[4,128,151936], param_1.1550: s32[4,128], param_2.1291: f32[4,128], param_3.911: f32[4,128], param_4.546: bf16[4,128], param_5.446: f32[4,128]) -> bf16[4,128,151936] { + %param_5.446 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.2087 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.446), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2078 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1355 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1448 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1355), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.544 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.544), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1448, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %mul.2086 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1356 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1454 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1356), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.546 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1454, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.62 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2077 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2078, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2077, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1543), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.2085 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2086, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1291 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1291), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2085, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1550 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.48 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.47 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1447 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1447), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2076 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2079, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1446 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2076), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.281 (param_0.1380: bf16[151936,2048], param_1.1562: f32[4,128], param_2.1324: bf16[4,128,2048], param_3.931: bf16[2048], param_4.567: bf16[4,128,151936], param_5.480: s32[4,128], param_6.369: f32[4,128], param_7.209: f32[4,128], param_8.126: bf16[4,128], param_9.95: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { - %param_4.567 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) - %param_5.480 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.369 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.209 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) - %param_8.126 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) - %param_9.95 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.567, %param_5.480, %param_6.369, %param_7.209, %param_8.126, /*index=5*/%param_9.95), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1562 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1324 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1453 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1453), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2084 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2087, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1452 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} +} + +%fused_computation.281 (param_0.1381: bf16[151936,2048], param_1.1569: f32[4,128], param_2.1327: bf16[4,128,2048], param_3.931: bf16[2048], param_4.569: bf16[4,128,151936], param_5.481: s32[4,128], param_6.371: f32[4,128], param_7.214: f32[4,128], param_8.131: bf16[4,128], param_9.98: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { + %param_4.569 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) + %param_5.481 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.371 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.214 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) + %param_8.131 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) + %param_9.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) + %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.569, %param_5.481, %param_6.371, %param_7.214, %param_8.131, /*index=5*/%param_9.98), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1327 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) %param_3.931 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.268.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1562, %param_2.1324, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.268.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1569, %param_2.1327, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.269.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %bitcast.333 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%convolution.86.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} %convert_element_type.1323 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_0.1380 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1380), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %param_0.1381 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} %add_any.184 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1323, %convert_element_type.1322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} %square.215 = f32[151936,2048]{1,0:T(8,128)} multiply(%add_any.184, %add_any.184), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1211 = f32[]{:T(128)} constant(0) - %reduce.177 = f32[]{:T(128)} reduce(%square.215, %constant.1211), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1213 = f32[]{:T(128)} constant(0) + %reduce.177 = f32[]{:T(128)} reduce(%square.215, %constant.1213), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.166 = (f32[]{:T(128)}, bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.177, %convolution.86.clone.1) } @@ -178,23 +178,23 @@ StackFrames ROOT %reduce_sum.394 = f32[]{:T(128)} add(%reduce_sum.389, %reduce_sum.393), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.288 (param_0.1391: bf16[4,128,151936], param_1.1570: f32[4,128], param_2.1327: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { - %param_2.1327 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1327), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} +%fused_computation.288 (param_0.1392: bf16[4,128,151936], param_1.1577: f32[4,128], param_2.1330: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { + %param_2.1330 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1330), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.25 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.24 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1391 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} %param_3.933 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) %sub.73 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.933), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.64 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1340, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1570 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1570), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1577 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1577), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.60 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.1223 = f32[]{:T(128)} constant(0) - %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.109"} + %constant.1225 = f32[]{:T(128)} constant(0) + %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1225), dimensions={}, metadata={op_name="broadcast.109"} %mul.1765 = f32[4,128,151936]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.769), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1765, %constant.1223), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1765, %constant.1225), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_9.12 (reduce_sum.186: f32[], reduce_sum.190: f32[]) -> f32[] { @@ -203,15 +203,15 @@ StackFrames ROOT %reduce_sum.191 = f32[]{:T(128)} add(%reduce_sum.186, %reduce_sum.190), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.293 (param_0.1392: bf16[4,128,151936], param_1.1571: bf16[4,128]) -> f32[4,128] { - %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1571 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1571), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} +%fused_computation.293 (param_0.1393: bf16[4,128,151936], param_1.1578: bf16[4,128]) -> f32[4,128] { + %param_0.1393 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_1.1578 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1578), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %sub.70 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1346, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.54 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.1224 = f32[]{:T(128)} constant(0) - ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1224), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %constant.1226 = f32[]{:T(128)} constant(0) + ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1226), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_33.38 (reduce_sum.269: f32[], reduce_sum.270: f32[]) -> f32[] { @@ -220,12 +220,12 @@ StackFrames ROOT %reduce_sum.274 = f32[]{:T(128)} add(%reduce_sum.269, %reduce_sum.270), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.298 (param_0.1386: f32[4,6144,2048]) -> f32[] { - %param_0.1386 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) - %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.298 (param_0.1387: f32[4,6144,2048]) -> f32[] { + %param_0.1387 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) + %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.218 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%bitcast.347, %bitcast.347), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1217 = f32[]{:T(128)} constant(0) - ROOT %reduce.181 = f32[]{:T(128)} reduce(%square.218, %constant.1217), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1219 = f32[]{:T(128)} constant(0) + ROOT %reduce.181 = f32[]{:T(128)} reduce(%square.218, %constant.1219), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_32.37 (reduce_sum.263: f32[], reduce_sum.267: f32[]) -> f32[] { @@ -240,35 +240,35 @@ StackFrames ROOT %reduce_sum.262 = f32[]{:T(128)} add(%reduce_sum.260, %reduce_sum.261), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.300 (param_0.1387: f32[4,2048,6144], param_1.1566: f32[4,2048,6144]) -> (f32[], f32[]) { - %param_0.1387 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) - %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.300 (param_0.1388: f32[4,2048,6144], param_1.1573: f32[4,2048,6144]) -> (f32[], f32[]) { + %param_0.1388 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) + %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.221 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.351, %bitcast.351), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1218 = f32[]{:T(128)} constant(0) - %reduce.182 = f32[]{:T(128)} reduce(%square.221, %constant.1218), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1566 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) - %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1566), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %constant.1220 = f32[]{:T(128)} constant(0) + %reduce.182 = f32[]{:T(128)} reduce(%square.221, %constant.1220), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1573 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) + %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.224.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.355.clone.1, %bitcast.355.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.183.clone.1 = f32[]{:T(128)} reduce(%square.224.clone.1, %constant.1218), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.183.clone.1 = f32[]{:T(128)} reduce(%square.224.clone.1, %constant.1220), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.167 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.182, %reduce.183.clone.1) } -%fused_computation.303 (param_0.900: f32[6144,4,2048]) -> bf16[4,6144,2048] { - %param_0.900 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) - %copy.192 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.900), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.303 (param_0.901: f32[6144,4,2048]) -> bf16[4,6144,2048] { + %param_0.901 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) + %copy.190 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.901), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} + ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.190), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.304 (param_0.902: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.902 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.193 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.902), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.304 (param_0.903: f32[2048,4,6144]) -> bf16[4,2048,6144] { + %param_0.903 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.191 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.903), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} + ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } -%fused_computation.305 (param_0.904: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.904 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.194 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.904), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.305 (param_0.905: f32[2048,4,6144]) -> bf16[4,2048,6144] { + %param_0.905 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.192 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.905), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} + ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_62.67 (reduce_sum.416: f32[], reduce_sum.417: f32[]) -> f32[] { @@ -283,39 +283,39 @@ StackFrames ROOT %reduce_sum.340 = f32[]{:T(128)} add(%reduce_sum.338, %reduce_sum.339), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.306 (param_0.1376: f32[6144,4,2048], param_1.1558: f32[], param_2.1320: f32[], param_3.927: f32[], param_4.563: f32[6144,4,2048], param_5.476: f32[], param_6.365: f32[4,6144,2048], param_7.205: pred[], param_8.122: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { - %param_0.1376 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) +%fused_computation.306 (param_0.1377: f32[6144,4,2048], param_1.1565: f32[], param_2.1323: f32[], param_3.927: f32[], param_4.565: f32[6144,4,2048], param_5.477: f32[], param_6.367: f32[4,6144,2048], param_7.210: pred[], param_8.127: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { + %param_0.1377 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) %param_3.927 = f32[]{:T(128)S(6)} parameter(3) %mul.1998.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_3.927), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.205 = pred[]{:T(512)S(6)} parameter(7) - %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.365 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) - %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.476 = f32[]{:T(128)} parameter(5) - %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.210 = pred[]{:T(512)S(6)} parameter(7) + %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.210), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.367 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) + %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.477 = f32[]{:T(128)} parameter(5) + %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.931.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%bitcast.482.clone.1, %div.932.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.303.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} select(%select_n.304.clone.1, %bitcast.482.clone.1, %div.931.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1146.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.886.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1146.clone.1), dimensions={}, metadata={op_name="broadcast.83"} %mul.2004.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.122 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) + %param_8.127 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) %constant.1150.clone.1 = f32[]{:T(128)} constant(0.9) %mul.2005.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1150.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2003.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.122, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2003.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.127, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.989.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2004.clone.1, %mul.2003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) - %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) + %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.74.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %select_n.303.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1149.clone.1 = f32[]{:T(128)} constant(0.05) %mul.2002.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1149.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.2000.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%integer_pow.74.clone.1, %mul.2002.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.563 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) + %param_4.565 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) %constant.1148.clone.1 = f32[]{:T(128)} constant(0.95) %mul.2001.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1148.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1999.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.563, %mul.2001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1999.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.565, %mul.2001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.988.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2000.clone.1, %mul.1999.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) - %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1565 = f32[]{:T(128)S(6)} parameter(1) + %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1565), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.926.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.988.clone.1, %div.927.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.71.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} sqrt(%div.926.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1147.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -323,14 +323,14 @@ StackFrames %add.986.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%sqrt.71.clone.1, %add.987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.435.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%div.928.clone.1, %add.986.clone.1), metadata={op_name="multiply.52"} %div.925.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.989.clone.1, %multiply.435.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1997.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1376, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1997.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.985.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%div.925.clone.1, %mul.1997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1996.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%mul.1998.clone.1, %add.985.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1376, %mul.1996.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1377, %mul.1996.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.225 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%add.984.clone.1, %add.984.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1207 = f32[]{:T(128)} constant(0) - %reduce.184 = f32[]{:T(128)} reduce(%square.225, %constant.1207), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1207), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1209 = f32[]{:T(128)} constant(0) + %reduce.184 = f32[]{:T(128)} reduce(%square.225, %constant.1209), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.145 = (f32[]{:T(128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.184, %add.984.clone.1, %add.988.clone.1, %add.989.clone.1, %reduce.187.clone.1) } @@ -346,39 +346,39 @@ StackFrames ROOT %reduce_sum.337 = f32[]{:T(128)} add(%reduce_sum.332, %reduce_sum.333), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.307 (param_0.1377: f32[2048,4,6144], param_1.1559: f32[], param_2.1321: f32[], param_3.928: f32[], param_4.564: f32[2048,4,6144], param_5.477: f32[], param_6.366: f32[4,2048,6144], param_7.206: pred[], param_8.123: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1377 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) +%fused_computation.307 (param_0.1378: f32[2048,4,6144], param_1.1566: f32[], param_2.1324: f32[], param_3.928: f32[], param_4.566: f32[2048,4,6144], param_5.478: f32[], param_6.368: f32[4,2048,6144], param_7.211: pred[], param_8.128: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { + %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) %param_3.928 = f32[]{:T(128)S(6)} parameter(3) %mul.2008.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.928), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.206 = pred[]{:T(512)S(6)} parameter(7) - %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.366 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.477 = f32[]{:T(128)} parameter(5) - %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.211 = pred[]{:T(512)S(6)} parameter(7) + %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.211), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.368 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.478 = f32[]{:T(128)} parameter(5) + %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.939.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.484.clone.1, %div.940.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.307.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.308.clone.1, %bitcast.484.clone.1, %div.939.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1152.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.892.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1152.clone.1), dimensions={}, metadata={op_name="broadcast.85"} %mul.2012.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.123 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) + %param_8.128 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) %constant.1156.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.891.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1156.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2011.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.123, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2011.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.128, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.994.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2012.clone.1, %mul.2011.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) - %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1324 = f32[]{:T(128)S(6)} parameter(2) + %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1324), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.75.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %select_n.307.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1155.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.890.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1155.clone.1), dimensions={}, metadata={op_name="broadcast.73"} %mul.2010.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.75.clone.1, %broadcast.890.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.564 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) + %param_4.566 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) %constant.1154.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.889.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1154.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2009.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.564, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2009.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.566, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.993.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2010.clone.1, %mul.2009.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) - %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1566 = f32[]{:T(128)S(6)} parameter(1) + %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1566), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.934.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.993.clone.1, %div.935.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.72.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.934.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1153.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -386,14 +386,14 @@ StackFrames %add.992.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.72.clone.1, %broadcast.887.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.436.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.936.clone.1, %add.992.clone.1), metadata={op_name="multiply.51"} %div.933.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.994.clone.1, %multiply.436.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2007.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2007.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.991.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.933.clone.1, %mul.2007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.2006.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2008.clone.1, %add.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1377, %mul.2006.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2006.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.226 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.990.clone.1, %add.990.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1208 = f32[]{:T(128)} constant(0) - %reduce.185 = f32[]{:T(128)} reduce(%square.226, %constant.1208), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1208), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1210 = f32[]{:T(128)} constant(0) + %reduce.185 = f32[]{:T(128)} reduce(%square.226, %constant.1210), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1210), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.146 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.185, %add.990.clone.1, %add.993.clone.1, %add.994.clone.1, %reduce.188.clone.1) } @@ -409,39 +409,39 @@ StackFrames ROOT %reduce_sum.331 = f32[]{:T(128)} add(%reduce_sum.326, %reduce_sum.330), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.308 (param_0.1378: f32[2048,4,6144], param_1.1560: f32[], param_2.1322: f32[], param_3.929: f32[], param_4.565: f32[2048,4,6144], param_5.478: f32[], param_6.367: f32[4,2048,6144], param_7.207: pred[], param_8.124: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) +%fused_computation.308 (param_0.1379: f32[2048,4,6144], param_1.1567: f32[], param_2.1325: f32[], param_3.929: f32[], param_4.567: f32[2048,4,6144], param_5.479: f32[], param_6.369: f32[4,2048,6144], param_7.212: pred[], param_8.129: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { + %param_0.1379 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) %param_3.929 = f32[]{:T(128)S(6)} parameter(3) %mul.2015.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.929), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.207 = pred[]{:T(512)S(6)} parameter(7) - %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.367 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.478 = f32[]{:T(128)} parameter(5) - %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.212 = pred[]{:T(512)S(6)} parameter(7) + %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.212), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.369 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.369), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.479 = f32[]{:T(128)} parameter(5) + %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.947.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.486.clone.1, %div.948.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.311.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.312.clone.1, %bitcast.486.clone.1, %div.947.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1158.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.898.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1158.clone.1), dimensions={}, metadata={op_name="broadcast.85"} %mul.2019.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.124 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) + %param_8.129 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) %constant.1162.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.897.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1162.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2018.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.124, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2018.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.129, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.999.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2019.clone.1, %mul.2018.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) - %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1325 = f32[]{:T(128)S(6)} parameter(2) + %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1325), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.76.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %select_n.311.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1161.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.896.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1161.clone.1), dimensions={}, metadata={op_name="broadcast.73"} %mul.2017.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.76.clone.1, %broadcast.896.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.565 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) + %param_4.567 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) %constant.1160.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.895.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1160.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2016.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.565, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2016.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.567, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.998.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2017.clone.1, %mul.2016.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) - %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1567 = f32[]{:T(128)S(6)} parameter(1) + %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1567), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.942.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.998.clone.1, %div.943.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.73.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.942.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1159.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -449,14 +449,14 @@ StackFrames %add.997.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.73.clone.1, %broadcast.893.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.437.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.944.clone.1, %add.997.clone.1), metadata={op_name="multiply.50"} %div.941.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.999.clone.1, %multiply.437.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2014.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2014.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1379, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.996.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.941.clone.1, %mul.2014.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.2013.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2015.clone.1, %add.996.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1379, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.227 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.995.clone.1, %add.995.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1209 = f32[]{:T(128)} constant(0) - %reduce.186 = f32[]{:T(128)} reduce(%square.227, %constant.1209), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1211 = f32[]{:T(128)} constant(0) + %reduce.186 = f32[]{:T(128)} reduce(%square.227, %constant.1211), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1211), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.147 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.186, %add.995.clone.1, %add.998.clone.1, %add.999.clone.1, %reduce.189.clone.1) } @@ -466,12 +466,12 @@ StackFrames ROOT %reduce_sum.304 = f32[]{:T(128)} add(%reduce_sum.302, %reduce_sum.303), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.324 (param_0.1381: f32[4,2048,16,128]) -> f32[] { - %param_0.1381 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.324 (param_0.1382: f32[4,2048,16,128]) -> f32[] { + %param_0.1382 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.230 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%bitcast.362, %bitcast.362), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1212 = f32[]{:T(128)} constant(0) - ROOT %reduce.190 = f32[]{:T(128)} reduce(%square.230, %constant.1212), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1214 = f32[]{:T(128)} constant(0) + ROOT %reduce.190 = f32[]{:T(128)} reduce(%square.230, %constant.1214), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_38.43 (reduce_sum.296: f32[], reduce_sum.297: f32[]) -> f32[] { @@ -480,18 +480,18 @@ StackFrames ROOT %reduce_sum.298 = f32[]{:T(128)} add(%reduce_sum.296, %reduce_sum.297), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.326 (param_0.1382: f32[4,16,128,2048]) -> f32[] { - %param_0.1382 = f32[4,16,128,2048]{3,2,0,1:T(8,128)S(1)} parameter(0) - %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.326 (param_0.1383: f32[4,16,128,2048]) -> f32[] { + %param_0.1383 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.233 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%bitcast.366, %bitcast.366), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1213 = f32[]{:T(128)} constant(0) - ROOT %reduce.191 = f32[]{:T(128)} reduce(%square.233, %constant.1213), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1215 = f32[]{:T(128)} constant(0) + ROOT %reduce.191 = f32[]{:T(128)} reduce(%square.233, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.327 (param_0.949: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { - %param_0.949 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) - %copy.195 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.949), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.195), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.327 (param_0.950: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { + %param_0.950 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) + %copy.193 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.950), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} + ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_68.73 (reduce_sum.449: f32[], reduce_sum.450: f32[]) -> f32[] { @@ -506,39 +506,39 @@ StackFrames ROOT %reduce_sum.373 = f32[]{:T(128)} add(%reduce_sum.368, %reduce_sum.372), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.328 (param_0.1370: f32[2048,4,16,128], param_1.1552: f32[], param_2.1314: f32[], param_3.921: f32[], param_4.557: f32[2048,4,16,128], param_5.470: f32[], param_6.359: f32[4,2048,16,128], param_7.199: pred[], param_8.116: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { - %param_0.1370 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.328 (param_0.1371: f32[2048,4,16,128], param_1.1559: f32[], param_2.1317: f32[], param_3.921: f32[], param_4.559: f32[2048,4,16,128], param_5.471: f32[], param_6.361: f32[4,2048,16,128], param_7.204: pred[], param_8.121: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { + %param_0.1371 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.921 = f32[]{:T(128)S(6)} parameter(3) %mul.1950.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_3.921), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.199 = pred[]{:T(512)S(6)} parameter(7) - %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.199), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.359 = f32[4,2048,16,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.470 = f32[]{:T(128)} parameter(5) - %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.204 = pred[]{:T(512)S(6)} parameter(7) + %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.361 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.471 = f32[]{:T(128)} parameter(5) + %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.883.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%bitcast.470.clone.1, %div.884.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.279.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} select(%select_n.280.clone.1, %bitcast.470.clone.1, %div.883.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1110.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.858.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1110.clone.1), dimensions={}, metadata={op_name="broadcast.75"} %mul.1956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.116 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) + %param_8.121 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1114.clone.1 = f32[]{:T(128)} constant(0.9) %mul.1957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1114.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.116, %mul.1957.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.121, %mul.1957.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1956.clone.1, %mul.1955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) - %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) + %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.68.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %select_n.279.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1113.clone.1 = f32[]{:T(128)} constant(0.05) %mul.1954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1113.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.1952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.68.clone.1, %mul.1954.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.557 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) + %param_4.559 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1112.clone.1 = f32[]{:T(128)} constant(0.95) %mul.1953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1112.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1951.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %mul.1953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1951.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.559, %mul.1953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1952.clone.1, %mul.1951.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1552 = f32[]{:T(128)S(6)} parameter(1) - %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1552), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) + %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.878.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.956.clone.1, %div.879.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.65.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} sqrt(%div.878.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1111.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -546,14 +546,14 @@ StackFrames %add.954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%sqrt.65.clone.1, %add.955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.429.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%div.880.clone.1, %add.954.clone.1), metadata={op_name="multiply.58"} %div.877.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.957.clone.1, %multiply.429.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1949.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1370, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1949.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%div.877.clone.1, %mul.1949.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1948.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%mul.1950.clone.1, %add.953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1370, %mul.1948.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.1948.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.234 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%add.952.clone.1, %add.952.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1201 = f32[]{:T(128)} constant(0) - %reduce.192 = f32[]{:T(128)} reduce(%square.234, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1203 = f32[]{:T(128)} constant(0) + %reduce.192 = f32[]{:T(128)} reduce(%square.234, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.148 = (f32[]{:T(128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.192, %add.952.clone.1, %add.956.clone.1, %add.957.clone.1, %reduce.194.clone.1) } @@ -569,39 +569,39 @@ StackFrames ROOT %reduce_sum.367 = f32[]{:T(128)} add(%reduce_sum.365, %reduce_sum.366), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.329 (param_0.1371: f32[16,4,128,2048], param_1.1553: f32[], param_2.1315: f32[], param_3.922: f32[], param_4.558: f32[16,4,128,2048], param_5.471: f32[], param_6.360: f32[4,16,128,2048], param_7.200: pred[], param_8.117: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { - %param_0.1371 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.329 (param_0.1372: f32[16,4,128,2048], param_1.1560: f32[], param_2.1318: f32[], param_3.922: f32[], param_4.560: f32[16,4,128,2048], param_5.472: f32[], param_6.362: f32[4,16,128,2048], param_7.205: pred[], param_8.122: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { + %param_0.1372 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) %param_3.922 = f32[]{:T(128)S(6)} parameter(3) %mul.1960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_3.922), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.200 = pred[]{:T(512)S(6)} parameter(7) - %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.200), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.360 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.471 = f32[]{:T(128)} parameter(5) - %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.205 = pred[]{:T(512)S(6)} parameter(7) + %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.362 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.472 = f32[]{:T(128)} parameter(5) + %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.891.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%bitcast.472.clone.1, %div.892.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.283.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} select(%select_n.284.clone.1, %bitcast.472.clone.1, %div.891.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1116.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.860.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1116.clone.1), dimensions={}, metadata={op_name="broadcast.76"} %mul.1966.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.117 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) + %param_8.122 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) %constant.1120.clone.1 = f32[]{:T(128)} constant(0.9) %mul.1967.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1120.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1965.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.117, %mul.1967.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1965.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.122, %mul.1967.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1966.clone.1, %mul.1965.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) - %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) + %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.69.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %select_n.283.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1119.clone.1 = f32[]{:T(128)} constant(0.05) %mul.1964.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1119.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.1962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%integer_pow.69.clone.1, %mul.1964.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.558 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) + %param_4.560 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) %constant.1118.clone.1 = f32[]{:T(128)} constant(0.95) %mul.1963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1118.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.558, %mul.1963.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.560, %mul.1963.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1962.clone.1, %mul.1961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1553 = f32[]{:T(128)S(6)} parameter(1) - %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1553), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) + %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.886.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.962.clone.1, %div.887.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.66.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} sqrt(%div.886.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1117.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -609,14 +609,14 @@ StackFrames %add.960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%sqrt.66.clone.1, %add.961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.430.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%div.888.clone.1, %add.960.clone.1), metadata={op_name="multiply.57"} %div.885.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.963.clone.1, %multiply.430.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1372, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%div.885.clone.1, %mul.1959.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%mul.1960.clone.1, %add.959.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.1958.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1372, %mul.1958.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.235 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%add.958.clone.1, %add.958.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1202 = f32[]{:T(128)} constant(0) - %reduce.193 = f32[]{:T(128)} reduce(%square.235, %constant.1202), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1202), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1204 = f32[]{:T(128)} constant(0) + %reduce.193 = f32[]{:T(128)} reduce(%square.235, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.149 = (f32[]{:T(128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.193, %add.958.clone.1, %add.962.clone.1, %add.963.clone.1, %reduce.195.clone.1) } @@ -632,23 +632,23 @@ StackFrames ROOT %reduce_sum.289 = f32[]{:T(128)} add(%reduce_sum.284, %reduce_sum.288), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.341 (param_0.1384: f32[4,2048,8,128], param_1.1564: f32[4,2048,8,128]) -> (f32[], f32[]) { - %param_0.1384 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) - %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.341 (param_0.1385: f32[4,2048,8,128], param_1.1571: f32[4,2048,8,128]) -> (f32[], f32[]) { + %param_0.1385 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.238 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1215 = f32[]{:T(128)} constant(0) - %reduce.196 = f32[]{:T(128)} reduce(%square.238, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1564 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) - %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %constant.1217 = f32[]{:T(128)} constant(0) + %reduce.196 = f32[]{:T(128)} reduce(%square.238, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1571 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(1) + %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.241.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.197.clone.1 = f32[]{:T(128)} reduce(%square.241.clone.1, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.197.clone.1 = f32[]{:T(128)} reduce(%square.241.clone.1, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.168 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.196, %reduce.197.clone.1) } -%fused_computation.344 (param_0.981: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.981 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %copy.196 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.981), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.196), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.344 (param_0.982: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { + %param_0.982 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.194 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.982), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} + ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} } %region_70.75 (reduce_sum.458: f32[], reduce_sum.459: f32[]) -> f32[] { @@ -663,39 +663,39 @@ StackFrames ROOT %reduce_sum.382 = f32[]{:T(128)} add(%reduce_sum.380, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.345 (param_0.1368: f32[2048,4,8,128], param_1.1550: f32[], param_2.1312: f32[], param_3.919: f32[], param_4.555: f32[2048,4,8,128], param_5.468: f32[], param_6.357: f32[4,2048,8,128], param_7.197: pred[], param_8.114: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1368 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) +%fused_computation.345 (param_0.1369: f32[2048,4,8,128], param_1.1557: f32[], param_2.1315: f32[], param_3.919: f32[], param_4.557: f32[2048,4,8,128], param_5.469: f32[], param_6.359: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { + %param_0.1369 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.919 = f32[]{:T(128)S(6)} parameter(3) %mul.1936.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.919), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.197 = pred[]{:T(512)S(6)} parameter(7) - %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.357 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.357), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.468 = f32[]{:T(128)} parameter(5) - %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.202 = pred[]{:T(512)S(6)} parameter(7) + %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.359 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.469 = f32[]{:T(128)} parameter(5) + %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.867.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.466.clone.1, %div.868.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.271.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.272.clone.1, %bitcast.466.clone.1, %div.867.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1098.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.850.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1098.clone.1), dimensions={}, metadata={op_name="broadcast.80"} %mul.1940.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.114 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1102.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.849.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1102.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.1939.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.114, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1939.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.946.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1940.clone.1, %mul.1939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1312 = f32[]{:T(128)S(6)} parameter(2) - %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1312), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) + %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.66.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %select_n.271.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1101.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.848.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1101.clone.1), dimensions={}, metadata={op_name="broadcast.69"} %mul.1938.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.66.clone.1, %broadcast.848.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.555 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %param_4.557 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1100.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.847.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1100.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1937.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.555, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1937.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.945.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1938.clone.1, %mul.1937.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1550 = f32[]{:T(128)S(6)} parameter(1) - %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) + %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.862.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.945.clone.1, %div.863.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.63.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.862.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1099.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -703,14 +703,14 @@ StackFrames %add.944.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.845.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.427.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.864.clone.1, %add.944.clone.1), metadata={op_name="multiply.60"} %div.861.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.946.clone.1, %multiply.427.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1935.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1368, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1935.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1369, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.943.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.861.clone.1, %mul.1935.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1934.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1936.clone.1, %add.943.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1368, %mul.1934.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1369, %mul.1934.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.242 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.942.clone.1, %add.942.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1199 = f32[]{:T(128)} constant(0) - %reduce.198 = f32[]{:T(128)} reduce(%square.242, %constant.1199), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1199), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1201 = f32[]{:T(128)} constant(0) + %reduce.198 = f32[]{:T(128)} reduce(%square.242, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.150 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.198, %add.942.clone.1, %add.945.clone.1, %add.946.clone.1, %reduce.200.clone.1) } @@ -726,39 +726,39 @@ StackFrames ROOT %reduce_sum.358 = f32[]{:T(128)} add(%reduce_sum.353, %reduce_sum.354), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.346 (param_0.1373: f32[2048,4,8,128], param_1.1555: f32[], param_2.1317: f32[], param_3.924: f32[], param_4.560: f32[2048,4,8,128], param_5.473: f32[], param_6.362: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1373 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) +%fused_computation.346 (param_0.1374: f32[2048,4,8,128], param_1.1562: f32[], param_2.1320: f32[], param_3.924: f32[], param_4.562: f32[2048,4,8,128], param_5.474: f32[], param_6.364: f32[4,2048,8,128], param_7.207: pred[], param_8.124: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { + %param_0.1374 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) %param_3.924 = f32[]{:T(128)S(6)} parameter(3) %mul.1977.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.924), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.202 = pred[]{:T(512)S(6)} parameter(7) - %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.362 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.473 = f32[]{:T(128)} parameter(5) - %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.207 = pred[]{:T(512)S(6)} parameter(7) + %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.364 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.474 = f32[]{:T(128)} parameter(5) + %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.907.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.476.clone.1, %div.908.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.291.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.292.clone.1, %bitcast.476.clone.1, %div.907.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1128.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.872.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1128.clone.1), dimensions={}, metadata={op_name="broadcast.80"} %mul.1981.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %param_8.124 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) %constant.1132.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.871.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1132.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.1980.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1980.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.124, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.973.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1981.clone.1, %mul.1980.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) - %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) + %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.71.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %select_n.291.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1131.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.870.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1131.clone.1), dimensions={}, metadata={op_name="broadcast.69"} %mul.1979.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.71.clone.1, %broadcast.870.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.560 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %param_4.562 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) %constant.1130.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.869.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1130.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1978.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.560, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1978.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.562, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.972.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1979.clone.1, %mul.1978.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1555 = f32[]{:T(128)S(6)} parameter(1) - %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1555), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1562 = f32[]{:T(128)S(6)} parameter(1) + %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1562), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.902.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.972.clone.1, %div.903.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.68.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.902.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1129.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -766,30 +766,30 @@ StackFrames %add.971.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.68.clone.1, %broadcast.867.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.432.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.904.clone.1, %add.971.clone.1), metadata={op_name="multiply.55"} %div.901.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.973.clone.1, %multiply.432.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1976.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1373, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1976.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1374, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.970.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.901.clone.1, %mul.1976.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1975.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1977.clone.1, %add.970.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1373, %mul.1975.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1374, %mul.1975.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.243 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.969.clone.1, %add.969.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1204 = f32[]{:T(128)} constant(0) - %reduce.199 = f32[]{:T(128)} reduce(%square.243, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) + %constant.1206 = f32[]{:T(128)} constant(0) + %reduce.199 = f32[]{:T(128)} reduce(%square.243, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) } -%fused_computation.362 (param_0.1055: bf16[4,128,2048], param_1.1114: f32[4,128], param_2.829: f32[4,128], param_3.497: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { - %param_3.497 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) +%fused_computation.362 (param_0.1056: bf16[4,128,2048], param_1.1117: f32[4,128], param_2.830: f32[4,128], param_3.495: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { + %param_3.495 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) %param_4.296 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.451 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.441 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.497, %dot_general.451), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.441), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.829 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1851 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.829), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %dot_general.448 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.438 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.495, %dot_general.448), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.438), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_2.830 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %mul.1851 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.830), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %mul.1843 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1363, %mul.1851), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.1055 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1055), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1850 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1114), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_0.1056 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1056), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1117 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %mul.1850 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1117), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %mul.1849 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1374, %mul.1850), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %add_any.193 = f32[4,128,2048]{2,1,0:T(8,128)} add(%mul.1843, %mul.1849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} ROOT %convert_element_type.1361 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.193), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} @@ -801,12 +801,12 @@ StackFrames ROOT %reduce_sum.185 = f32[]{:T(128)} add(%reduce_sum.171, %reduce_sum.184), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.363 (param_0.1393: bf16[4,128,2048]) -> f32[4,128] { - %param_0.1393 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.363 (param_0.1394: bf16[4,128,2048]) -> f32[4,128] { + %param_0.1394 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1394), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %square.246 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1365, %convert_element_type.1365), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} - %constant.1225 = f32[]{:T(128)} constant(0) - ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.246, %constant.1225), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + %constant.1227 = f32[]{:T(128)} constant(0) + ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.246, %constant.1227), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_12.15 (reduce_sum.198: f32[], reduce_sum.199: f32[]) -> f32[] { @@ -815,17 +815,17 @@ StackFrames ROOT %reduce_sum.200 = f32[]{:T(128)} add(%reduce_sum.198, %reduce_sum.199), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.365 (param_0.1388: bf16[4,128,2048], param_1.1567: bf16[4,128,2048], param_2.1325: bf16[2048]) -> f32[4,128] { - %param_0.1388 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1388), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1325 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1325), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.440 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1567, %dot_general.450), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.440), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.365 (param_0.1389: bf16[4,128,2048], param_1.1574: bf16[4,128,2048], param_2.1328: bf16[2048]) -> f32[4,128] { + %param_0.1389 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1389), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_1.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1328 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.447 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1328), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1574, %dot_general.447), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.437), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %mul.1847 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1372, %convert_element_type.1371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %constant.1219 = f32[]{:T(128)} constant(0) - ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1847, %constant.1219), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + %constant.1221 = f32[]{:T(128)} constant(0) + ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1847, %constant.1221), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} } %region_10.13 (dot_general.190: bf16[], dot_general.191: bf16[]) -> bf16[] { @@ -834,51 +834,51 @@ StackFrames ROOT %add.419 = bf16[]{:T(256)} add(%dot_general.190, %dot_general.191), metadata={op_name="add.82"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.285.clone.clone (param_0.1350: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1350 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.530 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1350), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%fused_computation.285.clone.clone (param_0.1351: bf16[151936,2048]) -> bf16[151936,2048,1] { + %param_0.1351 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} } -%fused_computation.289.clone.1.clone.clone (param_0.1351: bf16[4,128,151936], param_1.1539: s32[4,128], param_2.1282: f32[4,128], param_3.906: f32[4,128], param_4.540: bf16[4,128], param_5.441: f32[4,128]) -> bf16[4,128,151936] { - %param_5.441 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2067 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.441), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.289.clone.1.clone.clone (param_0.1352: bf16[4,128,151936], param_1.1546: s32[4,128], param_2.1285: f32[4,128], param_3.906: f32[4,128], param_4.542: bf16[4,128], param_5.442: f32[4,128]) -> bf16[4,128,151936] { + %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %mul.2075 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.442), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %param_3.906 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2066 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1351 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1438 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.540 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.540), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1438, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %mul.2074 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1352 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1444 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1352), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_4.542 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) + %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.542), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1444, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} %exp.60 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2065 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2066, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1282 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1282), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2065, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1539 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %mul.2073 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2074, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1285 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1285), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2073, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_1.1546 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.42 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} %eq.41 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1437 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1437), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2064 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2067, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1436 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2064), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %convert_element_type.1443 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1443), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2072 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2075, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %convert_element_type.1442 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2072), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} } -%fused_computation.366 (param_0.1349: f32[4,128], param_1.1538: bf16[4,128,2048], param_2.1283: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.541: s32[4,128], param_5.442: f32[4,128], param_6.338: f32[4,128], param_7.194: bf16[4,128], param_8.111: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { +%fused_computation.366 (param_0.1350: f32[4,128], param_1.1545: bf16[4,128,2048], param_2.1286: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.543: s32[4,128], param_5.443: f32[4,128], param_6.340: f32[4,128], param_7.199: bf16[4,128], param_8.116: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { %param_3.907 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.541 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.338 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.194 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.111 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.541, %param_5.442, %param_6.338, %param_7.194, /*index=5*/%param_8.111), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1283 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) - %fusion.250.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1283), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.250.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %param_1.1538 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1349 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1862 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1349), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %param_4.543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) + %param_5.443 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) + %param_6.340 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) + %param_7.199 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) + %param_8.116 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) + %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.543, %param_5.443, %param_6.340, %param_7.199, /*index=5*/%param_8.116), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} + %param_2.1286 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) + %fusion.251.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1286), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.251.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %param_1.1545 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1545), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1350 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.1862 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1350), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %mul.1861 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1384, %mul.1862), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %convert_element_type.1383 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1861), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %multiply.420 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%convolution.84.clone.1, %convert_element_type.1383), metadata={op_name="multiply.362"} @@ -887,11 +887,11 @@ StackFrames ROOT %tuple.165 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.204, %convolution.84.clone.1) } -%fused_computation.374 (param_0.1087: f32[64], param_1.1147: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1147 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1147), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.1087 = f32[64]{0:T(128)S(1)} parameter(0) - %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1087), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} +%fused_computation.374 (param_0.1088: f32[64], param_1.1150: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1150 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1150), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} + %param_0.1088 = f32[64]{0:T(128)S(1)} parameter(0) + %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1088), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %div.717 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.720, %div.718), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.717), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} %convert_element_type.1392 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} @@ -900,19 +900,19 @@ StackFrames ROOT %tuple.158 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1392, %convert_element_type.1391.clone.1) } -%fused_computation.375 (param_0.1084: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1084 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.375 (param_0.1085: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.1085 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1042 = bf16[]{:T(256)} constant(-inf) - %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1084, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1084, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.42 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.46, %pad.45), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } -%fused_computation.376 (param_0.1086: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1086 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) +%fused_computation.376 (param_0.1087: bf16[4,128,1,64]) -> bf16[4,128,1,128] { + %param_0.1087 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) %constant.1041 = bf16[]{:T(256)} constant(-inf) - %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1086, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1086, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} + %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} ROOT %maximum.43 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.48, %pad.47), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} } @@ -928,16 +928,16 @@ StackFrames ROOT %reduce_sum.277 = f32[]{:T(128)} add(%reduce_sum.275, %reduce_sum.276), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.380 (param_0.1385: f32[4,2048], param_1.1565: f32[4,2048]) -> (f32[], f32[]) { - %param_0.1385 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.380 (param_0.1386: f32[4,2048], param_1.1572: f32[4,2048]) -> (f32[], f32[]) { + %param_0.1386 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.249 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.404, %bitcast.404), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1216 = f32[]{:T(128)} constant(0) - %reduce.205 = f32[]{:T(128)} reduce(%square.249, %constant.1216), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1565 = f32[4,2048]{1,0:T(4,128)} parameter(1) - %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %constant.1218 = f32[]{:T(128)} constant(0) + %reduce.205 = f32[]{:T(128)} reduce(%square.249, %constant.1218), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1572 = f32[4,2048]{1,0:T(4,128)} parameter(1) + %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.252.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.408.clone.1, %bitcast.408.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.206.clone.1 = f32[]{:T(128)} reduce(%square.252.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.206.clone.1 = f32[]{:T(128)} reduce(%square.252.clone.1, %constant.1218), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.169 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.205, %reduce.206.clone.1) } @@ -953,39 +953,39 @@ StackFrames ROOT %reduce_sum.352 = f32[]{:T(128)} add(%reduce_sum.347, %reduce_sum.351), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.383 (param_0.1374: f32[2048,4], param_1.1556: f32[], param_2.1318: f32[], param_3.925: f32[], param_4.561: f32[2048,4], param_5.474: f32[], param_6.363: f32[4,2048], param_7.203: pred[], param_8.120: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1374 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.383 (param_0.1375: f32[2048,4], param_1.1563: f32[], param_2.1321: f32[], param_3.925: f32[], param_4.563: f32[2048,4], param_5.475: f32[], param_6.365: f32[4,2048], param_7.208: pred[], param_8.125: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.925 = f32[]{:T(128)S(6)} parameter(3) %mul.1984.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.925), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.203 = pred[]{:T(512)S(6)} parameter(7) - %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.363 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.474 = f32[]{:T(128)} parameter(5) - %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.208 = pred[]{:T(512)S(6)} parameter(7) + %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.365 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.475 = f32[]{:T(128)} parameter(5) + %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.915.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.478.clone.1, %div.916.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.295.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.296.clone.1, %bitcast.478.clone.1, %div.915.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1134.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.878.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1134.clone.1), dimensions={}, metadata={op_name="broadcast.82"} %mul.1988.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.120 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %param_8.125 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1138.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.877.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1138.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.1987.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1987.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.125, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.978.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1988.clone.1, %mul.1987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) - %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) + %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.72.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %select_n.295.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1137.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.876.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1137.clone.1), dimensions={}, metadata={op_name="broadcast.71"} %mul.1986.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.72.clone.1, %broadcast.876.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.561 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %param_4.563 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1136.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.875.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1136.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1985.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1985.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.563, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.977.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1986.clone.1, %mul.1985.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) - %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1563 = f32[]{:T(128)S(6)} parameter(1) + %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1563), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.910.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.977.clone.1, %div.911.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.69.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.910.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1135.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -993,14 +993,14 @@ StackFrames %add.976.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.69.clone.1, %broadcast.873.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.433.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.912.clone.1, %add.976.clone.1), metadata={op_name="multiply.54"} %div.909.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.978.clone.1, %multiply.433.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1983.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1374, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1983.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.975.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.909.clone.1, %mul.1983.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1982.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1984.clone.1, %add.975.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1374, %mul.1982.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.1982.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.253 = f32[2048,4]{0,1:T(4,128)} multiply(%add.974.clone.1, %add.974.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1205 = f32[]{:T(128)} constant(0) - %reduce.207 = f32[]{:T(128)} reduce(%square.253, %constant.1205), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1207 = f32[]{:T(128)} constant(0) + %reduce.207 = f32[]{:T(128)} reduce(%square.253, %constant.1207), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1207), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.152 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.207, %add.974.clone.1, %add.977.clone.1, %add.978.clone.1, %reduce.209.clone.1) } @@ -1016,39 +1016,39 @@ StackFrames ROOT %reduce_sum.346 = f32[]{:T(128)} add(%reduce_sum.344, %reduce_sum.345), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.384 (param_0.1375: f32[2048,4], param_1.1557: f32[], param_2.1319: f32[], param_3.926: f32[], param_4.562: f32[2048,4], param_5.475: f32[], param_6.364: f32[4,2048], param_7.204: pred[], param_8.121: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.384 (param_0.1376: f32[2048,4], param_1.1564: f32[], param_2.1322: f32[], param_3.926: f32[], param_4.564: f32[2048,4], param_5.476: f32[], param_6.366: f32[4,2048], param_7.209: pred[], param_8.126: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1376 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.926 = f32[]{:T(128)S(6)} parameter(3) %mul.1991.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.926), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.204 = pred[]{:T(512)S(6)} parameter(7) - %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.364 = f32[4,2048]{1,0:T(4,128)} parameter(6) - %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.475 = f32[]{:T(128)} parameter(5) - %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.209 = pred[]{:T(512)S(6)} parameter(7) + %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.209), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.366 = f32[4,2048]{1,0:T(4,128)} parameter(6) + %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.476 = f32[]{:T(128)} parameter(5) + %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.923.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.480.clone.1, %div.924.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.299.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.300.clone.1, %bitcast.480.clone.1, %div.923.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1140.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.884.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1140.clone.1), dimensions={}, metadata={op_name="broadcast.82"} %mul.1995.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.121 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %param_8.126 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1144.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.883.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1144.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.1994.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.121, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1994.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.126, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.983.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1995.clone.1, %mul.1994.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) - %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) + %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.73.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %select_n.299.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1143.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.882.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1143.clone.1), dimensions={}, metadata={op_name="broadcast.71"} %mul.1993.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.73.clone.1, %broadcast.882.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.562 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %param_4.564 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1142.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.881.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1142.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1992.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.562, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1992.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.564, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.982.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1993.clone.1, %mul.1992.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) - %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1564 = f32[]{:T(128)S(6)} parameter(1) + %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1564), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.918.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.982.clone.1, %div.919.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.70.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.918.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1141.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1056,14 +1056,14 @@ StackFrames %add.981.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.70.clone.1, %broadcast.879.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.434.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.920.clone.1, %add.981.clone.1), metadata={op_name="multiply.53"} %div.917.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.983.clone.1, %multiply.434.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1990.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1990.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1376, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.980.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.917.clone.1, %mul.1990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1989.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1991.clone.1, %add.980.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.1989.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1376, %mul.1989.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.254 = f32[2048,4]{0,1:T(4,128)} multiply(%add.979.clone.1, %add.979.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1206 = f32[]{:T(128)} constant(0) - %reduce.208 = f32[]{:T(128)} reduce(%square.254, %constant.1206), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1206), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1208 = f32[]{:T(128)} constant(0) + %reduce.208 = f32[]{:T(128)} reduce(%square.254, %constant.1208), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1208), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.153 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.208, %add.979.clone.1, %add.982.clone.1, %add.983.clone.1, %reduce.210.clone.1) } @@ -1073,12 +1073,12 @@ StackFrames ROOT %reduce_sum.197 = f32[]{:T(128)} add(%reduce_sum.192, %reduce_sum.193), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.395 (param_0.1389: bf16[2048]) -> f32[] { - %param_0.1389 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.395 (param_0.1390: bf16[2048]) -> f32[] { + %param_0.1390 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} %square.257 = f32[2048]{0:T(1024)} multiply(%convert_element_type.1396, %convert_element_type.1396), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1220 = f32[]{:T(128)} constant(0) - ROOT %reduce.211 = f32[]{:T(128)} reduce(%square.257, %constant.1220), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1222 = f32[]{:T(128)} constant(0) + ROOT %reduce.211 = f32[]{:T(128)} reduce(%square.257, %constant.1222), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_59.64 (reduce_sum.401: f32[], reduce_sum.402: f32[]) -> f32[] { @@ -1093,39 +1093,39 @@ StackFrames ROOT %reduce_sum.325 = f32[]{:T(128)} add(%reduce_sum.323, %reduce_sum.324), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.396 (param_0.1379: f32[2048], param_1.1561: f32[], param_2.1323: f32[], param_3.930: f32[], param_4.566: f32[2048], param_5.479: f32[], param_6.368: bf16[2048], param_7.208: pred[], param_8.125: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { - %param_0.1379 = f32[2048]{0:T(1024)S(1)} parameter(0) +%fused_computation.396 (param_0.1380: f32[2048], param_1.1568: f32[], param_2.1326: f32[], param_3.930: f32[], param_4.568: f32[2048], param_5.480: f32[], param_6.370: bf16[2048], param_7.213: pred[], param_8.130: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { + %param_0.1380 = f32[2048]{0:T(1024)S(1)} parameter(0) %param_3.930 = f32[]{:T(128)S(6)} parameter(3) %mul.2022.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_3.930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.208 = pred[]{:T(512)S(6)} parameter(7) - %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.368 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.479 = f32[]{:T(128)} parameter(5) - %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.213 = pred[]{:T(512)S(6)} parameter(7) + %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.213), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.370 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_5.480 = f32[]{:T(128)} parameter(5) + %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.480), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.955.clone.1 = f32[2048]{0:T(1024)} divide(%convert_element_type.1411.clone.1, %div.956.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.315.clone.1 = f32[2048]{0:T(1024)} select(%select_n.316.clone.1, %convert_element_type.1411.clone.1, %div.955.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1164.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.900.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1164.clone.1), dimensions={}, metadata={op_name="broadcast.86"} %mul.2028.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.125 = f32[2048]{0:T(1024)S(1)} parameter(8) + %param_8.130 = f32[2048]{0:T(1024)S(1)} parameter(8) %constant.1168.clone.1 = f32[]{:T(128)} constant(0.9) %mul.2029.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1168.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2027.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.125, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2027.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.130, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.1005.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2028.clone.1, %mul.2027.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) - %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1326 = f32[]{:T(128)S(6)} parameter(2) + %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1326), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.77.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %select_n.315.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1167.clone.1 = f32[]{:T(128)} constant(0.05) %mul.2026.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1167.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %mul.2024.clone.1 = f32[2048]{0:T(1024)} multiply(%integer_pow.77.clone.1, %mul.2026.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.566 = f32[2048]{0:T(1024)S(1)} parameter(4) + %param_4.568 = f32[2048]{0:T(1024)S(1)} parameter(4) %constant.1166.clone.1 = f32[]{:T(128)} constant(0.95) %mul.2025.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1166.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2023.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.566, %mul.2025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2023.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.568, %mul.2025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.1004.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2024.clone.1, %mul.2023.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) - %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) + %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.950.clone.1 = f32[2048]{0:T(1024)} divide(%add.1004.clone.1, %div.951.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.74.clone.1 = f32[2048]{0:T(1024)} sqrt(%div.950.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1165.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1133,37 +1133,37 @@ StackFrames %add.1002.clone.1 = f32[2048]{0:T(1024)} add(%sqrt.74.clone.1, %add.1003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.438.clone.1 = f32[2048]{0:T(1024)} multiply(%div.952.clone.1, %add.1002.clone.1), metadata={op_name="multiply.49"} %div.949.clone.1 = f32[2048]{0:T(1024)} divide(%add.1005.clone.1, %multiply.438.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2021.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1379, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2021.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1380, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.1001.clone.1 = f32[2048]{0:T(1024)} add(%div.949.clone.1, %mul.2021.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.2020.clone.1 = f32[2048]{0:T(1024)} multiply(%mul.2022.clone.1, %add.1001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1379, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1380, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.258 = f32[2048]{0:T(1024)} multiply(%add.1000.clone.1, %add.1000.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1210 = f32[]{:T(128)} constant(0) - %reduce.212 = f32[]{:T(128)} reduce(%square.258, %constant.1210), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1210), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1212 = f32[]{:T(128)} constant(0) + %reduce.212 = f32[]{:T(128)} reduce(%square.258, %constant.1212), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1212), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.156 = (f32[]{:T(128)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.212, %add.1000.clone.1, %add.1004.clone.1, %add.1005.clone.1, %reduce.213.clone.1) } -%fused_computation.402 (param_0.1149: s32[512]) -> s32[1024] { +%fused_computation.402 (param_0.1150: s32[512]) -> s32[1024] { %constant.972 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.815 = s32[1024]{0:T(1024)} broadcast(%constant.972), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.1149 = s32[512]{0:T(512)S(1)} parameter(0) + %param_0.1150 = s32[512]{0:T(512)S(1)} parameter(0) %constant.973 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1149, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1150, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %constant.971 = s32[] constant(151935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} %broadcast.814 = s32[1024]{0:T(1024)} broadcast(%constant.971), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.815, %pad.49, %broadcast.814), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} } -%fused_computation.405 (param_0.1148: s32[4,128]) -> s32[512] { - %param_0.1148 = s32[4,128]{1,0:T(4,128)} parameter(0) +%fused_computation.405 (param_0.1149: s32[4,128]) -> s32[512] { + %param_0.1149 = s32[4,128]{1,0:T(4,128)} parameter(0) %constant.1065 = s32[]{:T(128)} constant(0) %broadcast.834 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1065), dimensions={}, metadata={op_name="broadcast.95"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1148, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} + %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1149, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} %constant.1051 = s32[]{:T(128)} constant(151936) %add.925 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1051), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1148, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1148), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} + %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1149, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} + %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1149), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} ROOT %bitcast.409 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} } @@ -1179,16 +1179,16 @@ StackFrames ROOT %reduce_sum.295 = f32[]{:T(128)} add(%reduce_sum.290, %reduce_sum.291), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.407 (param_0.1383: f32[4,128], param_1.1563: f32[4,128]) -> (f32[], f32[]) { - %param_0.1383 = f32[4,128]{1,0:T(4,128)} parameter(0) - %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} +%fused_computation.407 (param_0.1384: f32[4,128], param_1.1570: f32[4,128]) -> (f32[], f32[]) { + %param_0.1384 = f32[4,128]{1,0:T(4,128)} parameter(0) + %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.261 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.413, %bitcast.413), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1214 = f32[]{:T(128)} constant(0) - %reduce.214 = f32[]{:T(128)} reduce(%square.261, %constant.1214), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1563 = f32[4,128]{1,0:T(4,128)} parameter(1) - %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %constant.1216 = f32[]{:T(128)} constant(0) + %reduce.214 = f32[]{:T(128)} reduce(%square.261, %constant.1216), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1570 = f32[4,128]{1,0:T(4,128)} parameter(1) + %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} %square.264.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.417.clone.1, %bitcast.417.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.215.clone.1 = f32[]{:T(128)} reduce(%square.264.clone.1, %constant.1214), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.215.clone.1 = f32[]{:T(128)} reduce(%square.264.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.170 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.214, %reduce.215.clone.1) } @@ -1204,27 +1204,27 @@ StackFrames ROOT %reduce_sum.400 = f32[]{:T(128)} add(%reduce_sum.395, %reduce_sum.396), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.410 (param_0.1390: bf16[4,128], param_1.1569: f32[4,128], param_2.1326: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { +%fused_computation.410 (param_0.1391: bf16[4,128], param_1.1576: f32[4,128], param_2.1329: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { %param_3.932 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) %constant.1170.clone.1 = s32[]{:T(128)} constant(0) %broadcast.901.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1170.clone.1), dimensions={}, metadata={op_name="broadcast.95"} %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.932, %broadcast.901.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1569), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1390 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1390), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %param_1.1576 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1576), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_0.1391 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} %add.927 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} %square.269 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %add.927), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} - %constant.1222 = f32[]{:T(128)} constant(0) - %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1222), dimensions={}, metadata={op_name="broadcast.99"} + %constant.1224 = f32[]{:T(128)} constant(0) + %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1224), dimensions={}, metadata={op_name="broadcast.99"} %mul.1913 = f32[4,128]{1,0:T(4,128)} multiply(%square.269, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} %mul.1893 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1913, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.216 = f32[]{:T(128)} reduce(%mul.1893, %constant.1222), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1326 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1326), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %reduce.216 = f32[]{:T(128)} reduce(%mul.1893, %constant.1224), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_2.1329 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1329), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} %add.904.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1913), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} %mul.1894.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.904.clone.1, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1894.clone.1, %constant.1222), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1894.clone.1, %constant.1224), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} %mul.1911.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %broadcast.831), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} %constant.1068.clone.1 = f32[]{:T(128)} constant(1) %add.922.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1068.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} @@ -1244,39 +1244,39 @@ StackFrames ROOT %reduce_sum.379 = f32[]{:T(128)} add(%reduce_sum.374, %reduce_sum.375), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.411 (param_0.1369: f32[128,4], param_1.1551: f32[], param_2.1313: f32[], param_3.920: f32[], param_4.556: f32[128,4], param_5.469: f32[], param_6.358: f32[4,128], param_7.198: pred[], param_8.115: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1369 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.411 (param_0.1370: f32[128,4], param_1.1558: f32[], param_2.1316: f32[], param_3.920: f32[], param_4.558: f32[128,4], param_5.470: f32[], param_6.360: f32[4,128], param_7.203: pred[], param_8.120: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1370 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.920 = f32[]{:T(128)S(6)} parameter(3) %mul.1943.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.920), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.198 = pred[]{:T(512)S(6)} parameter(7) - %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.198), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.358 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.469 = f32[]{:T(128)} parameter(5) - %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.203 = pred[]{:T(512)S(6)} parameter(7) + %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.360 = f32[4,128]{1,0:T(4,128)} parameter(6) + %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.470 = f32[]{:T(128)} parameter(5) + %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.875.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.468.clone.1, %div.876.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.275.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.276.clone.1, %bitcast.468.clone.1, %div.875.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1104.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.856.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1104.clone.1), dimensions={}, metadata={op_name="broadcast.78"} %mul.1947.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.115 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %param_8.120 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1108.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.855.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1108.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.1946.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.115, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1946.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.951.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1947.clone.1, %mul.1946.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1313 = f32[]{:T(128)S(6)} parameter(2) - %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1313), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) + %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.67.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %select_n.275.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1107.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.854.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1107.clone.1), dimensions={}, metadata={op_name="broadcast.67"} %mul.1945.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.854.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.556 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %param_4.558 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1106.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.853.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1106.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1944.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.556, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1944.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.558, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.950.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1945.clone.1, %mul.1944.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1551 = f32[]{:T(128)S(6)} parameter(1) - %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1551), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) + %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.870.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.950.clone.1, %div.871.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.64.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.870.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1105.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1284,14 +1284,14 @@ StackFrames %add.949.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.851.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.428.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.872.clone.1, %add.949.clone.1), metadata={op_name="multiply.59"} %div.869.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.951.clone.1, %multiply.428.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1942.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1369, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1942.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1370, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.948.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.869.clone.1, %mul.1942.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1941.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1943.clone.1, %add.948.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1369, %mul.1941.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1370, %mul.1941.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.265 = f32[128,4]{0,1:T(4,128)} multiply(%add.947.clone.1, %add.947.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1200 = f32[]{:T(128)} constant(0) - %reduce.217 = f32[]{:T(128)} reduce(%square.265, %constant.1200), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1202 = f32[]{:T(128)} constant(0) + %reduce.217 = f32[]{:T(128)} reduce(%square.265, %constant.1202), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1202), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.159 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.217, %add.947.clone.1, %add.950.clone.1, %add.951.clone.1, %reduce.221.clone.1) } @@ -1307,39 +1307,39 @@ StackFrames ROOT %reduce_sum.361 = f32[]{:T(128)} add(%reduce_sum.359, %reduce_sum.360), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.412 (param_0.1372: f32[128,4], param_1.1554: f32[], param_2.1316: f32[], param_3.923: f32[], param_4.559: f32[128,4], param_5.472: f32[], param_6.361: f32[4,128], param_7.201: pred[], param_8.118: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1372 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) +%fused_computation.412 (param_0.1373: f32[128,4], param_1.1561: f32[], param_2.1319: f32[], param_3.923: f32[], param_4.561: f32[128,4], param_5.473: f32[], param_6.363: f32[4,128], param_7.206: pred[], param_8.123: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1373 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) %param_3.923 = f32[]{:T(128)S(6)} parameter(3) %mul.1970.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.923), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.201 = pred[]{:T(512)S(6)} parameter(7) - %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.201), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.361 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.472 = f32[]{:T(128)} parameter(5) - %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_7.206 = pred[]{:T(512)S(6)} parameter(7) + %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.363 = f32[4,128]{1,0:T(4,128)} parameter(6) + %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + %param_5.473 = f32[]{:T(128)} parameter(5) + %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.899.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.474.clone.1, %div.900.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %select_n.287.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.288.clone.1, %bitcast.474.clone.1, %div.899.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} %constant.1122.clone.1 = f32[]{:T(128)} constant(0.1) %broadcast.866.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1122.clone.1), dimensions={}, metadata={op_name="broadcast.78"} %mul.1974.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.118 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %param_8.123 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) %constant.1126.clone.1 = f32[]{:T(128)} constant(0.9) %broadcast.865.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1126.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.1973.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.118, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1973.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.123, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.968.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1974.clone.1, %mul.1973.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) - %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) + %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %integer_pow.70.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %select_n.287.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} %constant.1125.clone.1 = f32[]{:T(128)} constant(0.05) %broadcast.864.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1125.clone.1), dimensions={}, metadata={op_name="broadcast.67"} %mul.1972.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.864.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.559 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %param_4.561 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) %constant.1124.clone.1 = f32[]{:T(128)} constant(0.95) %broadcast.863.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1124.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1971.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.559, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1971.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.967.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1972.clone.1, %mul.1971.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1554 = f32[]{:T(128)S(6)} parameter(1) - %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1554), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) + %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} %div.894.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.967.clone.1, %div.895.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} %sqrt.67.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.894.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} %constant.1123.clone.1 = f32[]{:T(128)} constant(1e-08) @@ -1347,23 +1347,23 @@ StackFrames %add.966.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.67.clone.1, %broadcast.861.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %multiply.431.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.896.clone.1, %add.966.clone.1), metadata={op_name="multiply.56"} %div.893.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.968.clone.1, %multiply.431.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1969.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1372, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1969.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1373, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} %add.965.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.893.clone.1, %mul.1969.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %mul.1968.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1970.clone.1, %add.965.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1372, %mul.1968.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1373, %mul.1968.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} %square.266 = f32[128,4]{0,1:T(4,128)} multiply(%add.964.clone.1, %add.964.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1203 = f32[]{:T(128)} constant(0) - %reduce.218 = f32[]{:T(128)} reduce(%square.266, %constant.1203), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1203), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %constant.1205 = f32[]{:T(128)} constant(0) + %reduce.218 = f32[]{:T(128)} reduce(%square.266, %constant.1205), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} ROOT %tuple.160 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.218, %add.964.clone.1, %add.967.clone.1, %add.968.clone.1, %reduce.222.clone.1) } -%fused_computation.421 (param_0.1200: f32[4,128], param_1.1320: f32[4,128]) -> f32[4,128] { - %param_0.1200 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1320 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) +%fused_computation.421 (param_0.1201: f32[4,128], param_1.1323: f32[4,128]) -> f32[4,128] { + %param_0.1201 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %param_1.1323 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) %constant.1045 = f32[]{:T(128)} constant(0.00048828125) %broadcast.837 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1045), dimensions={}, metadata={op_name="broadcast.399"} - %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1320, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1323, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1043 = f32[]{:T(128)} constant(1e-06) %add.935 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1043), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.934 = f32[4,128]{1,0:T(4,128)} add(%div.767, %add.935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} @@ -1372,7 +1372,7 @@ StackFrames %constant.1040 = f32[]{:T(128)} constant(-0.5) %mul.1919 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1040), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %mul.1910 = f32[4,128]{1,0:T(4,128)} multiply(%div.754, %mul.1919), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1909 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1200, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.1909 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1201, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} %constant.1039 = f32[]{:T(128)} constant(0.0009765625) %mul.1918 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1039), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} ROOT %mul.1908 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1909, %mul.1918), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} @@ -1384,31 +1384,31 @@ StackFrames ROOT %reduce_sum.139 = s32[]{:T(128)} add(%reduce_sum.137, %reduce_sum.138), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.425 (param_0.1219: pred[4,128]) -> s32[] { - %param_0.1219 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1219), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} +%fused_computation.425 (param_0.1220: pred[4,128]) -> s32[] { + %param_0.1220 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1220), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} %constant.1066 = s32[]{:T(128)} constant(0) ROOT %reduce.220 = s32[]{:T(128)} reduce(%convert_element_type.1403, %constant.1066), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.428 (param_0.1202: f32[4,128]) -> f32[4,128] { - %param_0.1202 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) +%fused_computation.428 (param_0.1203: f32[4,128]) -> f32[4,128] { + %param_0.1203 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) %constant.1046 = f32[]{:T(128)} constant(0.00048828125) %broadcast.829 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1046), dimensions={}, metadata={op_name="broadcast.399"} - %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1202, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} + %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1203, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} %constant.1044 = f32[]{:T(128)} constant(1e-06) %add.924 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1044), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} %add.921 = f32[4,128]{1,0:T(4,128)} add(%div.759, %add.924), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} ROOT %rsqrt.166 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.921), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} } -%fused_computation.429 (param_0.1203: pred[4,128], param_1.1568: f32[]) -> f32[4,128] { - %param_0.1203 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} - %constant.1221 = f32[]{:T(128)} constant(0) - %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1221), dimensions={}, metadata={op_name="broadcast.99"} - ROOT %mul.1920 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1203, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.429 (param_0.1204: pred[4,128], param_1.1575: f32[]) -> f32[4,128] { + %param_0.1204 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) + %param_1.1575 = f32[]{:T(128)S(6)} parameter(1) + %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1575), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} + %constant.1223 = f32[]{:T(128)} constant(0) + %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.99"} + ROOT %mul.1920 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1204, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} } %fused_computation.431 () -> f32[64] { @@ -1425,35 +1425,35 @@ StackFrames ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.840, %div.768), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} } -%fused_computation.432 (param_0.1217: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.1217 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1217), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} +%fused_computation.432 (param_0.1218: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { + %param_0.1218 = s32[4,128]{1,0:T(4,128)} parameter(0) + %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1218), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} %bitcast.418 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} ROOT %tuple.162 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.418, %convert_element_type.1405) } -%fused_computation.435 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { - %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.533 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.533) +%fused_computation.435 (param_0.1360: f32[2048,4]) -> bf16[4,2048] { + %param_0.1360 = f32[2048,4]{0,1:T(4,128)} parameter(0) + %bitcast.531 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.531) } -%fused_computation.436 (param_0.1358: f32[2048,4]) -> bf16[4,2048] { - %param_0.1358 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.532 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.532) +%fused_computation.436 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { + %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) + %bitcast.530 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.530) } -%fused_computation.437 (param_0.1360: f32[128,4]) -> bf16[4,128] { - %param_0.1360 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.534 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.534) +%fused_computation.437 (param_0.1361: f32[128,4]) -> bf16[4,128] { + %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) + %bitcast.532 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.532) } -%fused_computation.438 (param_0.1361: f32[128,4]) -> bf16[4,128] { - %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.535 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.535) +%fused_computation.438 (param_0.1362: f32[128,4]) -> bf16[4,128] { + %param_0.1362 = f32[128,4]{0,1:T(4,128)} parameter(0) + %bitcast.533 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1362), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.533) } %region_8.11 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { @@ -1462,40 +1462,40 @@ StackFrames ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.287.clone.clone (param_0.1345: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1345 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1345), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%fused_computation.287.clone.clone (param_0.1346: bf16[151936,2048]) -> bf16[151936,2048,1] { + %param_0.1346 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.526 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1346), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} } -%fused_computation.368.clone.clone (param_0.1346: f32[4,128], param_1.1535: bf16[4,128,2048], param_2.1278: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1278 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1278), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1535 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1432 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1535), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1346 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2059 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1346), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2058 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1432, %mul.2059), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1431 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2058), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.474 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.475, %convert_element_type.1431), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.368.clone.clone (param_0.1347: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1281: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1281 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.476 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1281), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1438 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + %param_0.1347 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2067 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1347), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %mul.2066 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1438, %mul.2067), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} + %convert_element_type.1437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2066), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} + ROOT %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.476, %convert_element_type.1437), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.439 (param_0.1362: bf16[151936,2048], param_1.1544: f32[4,128], param_2.1302: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { - %param_1.1544 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1302 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) +%fused_computation.439 (param_0.1363: bf16[151936,2048], param_1.1551: f32[4,128], param_2.1305: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { + %param_1.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1305 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) %param_3.913 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1544, %param_2.1302, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1362 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %fusion.252.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1362), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.269.clone.1, %fusion.252.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %constant.1193 = bf16[]{:T(256)} constant(-inf) - %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1193), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %fusion.270.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1551, %param_2.1305, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1363 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %fusion.253.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1363), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} + %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.270.clone.1, %fusion.253.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} + %constant.1195 = bf16[]{:T(256)} constant(-inf) + %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1195), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} ROOT %tuple.164 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,151936]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.223, %convolution.85.clone.1) } -%fused_computation.440 (param_0.1357: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.1357 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %bitcast.531 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1357), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.531) +%fused_computation.440 (param_0.1358: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { + %param_0.1358 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %bitcast.529 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} + ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.529) } %convert_element_type.767.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { @@ -1504,13 +1504,13 @@ StackFrames ROOT %add.755 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.155.clone.clone (param_0.1533: bf16[4,2048], param_1.1680: s32[]) -> bf16[2048] { - %param_0.1533 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1680 = s32[]{:T(128)S(6)} parameter(1) - %constant.1359 = s32[]{:T(128)} constant(0) - %dynamic_slice.384 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1533, %param_1.1680, %constant.1359), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1360 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.384, %constant.1360), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.155.clone.clone (param_0.1534: bf16[4,2048], param_1.1687: s32[]) -> bf16[2048] { + %param_0.1534 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) + %constant.1361 = s32[]{:T(128)} constant(0) + %dynamic_slice.388 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1687, %constant.1361), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1362 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.388, %constant.1362), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %region_14.16 (reduce_sum.204: f32[], reduce_sum.205: f32[]) -> f32[] { @@ -1519,25 +1519,25 @@ StackFrames ROOT %reduce_sum.206 = f32[]{:T(128)} add(%reduce_sum.204, %reduce_sum.205), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.58.clone.clone (param_0.1534: bf16[4,4,128,2048], param_1.1681: s32[]) -> f32[4,128] { - %param_0.1534 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1681 = s32[]{:T(128)S(6)} parameter(1) - %constant.1361 = s32[]{:T(128)} constant(0) - %dynamic_slice.385 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1681, %constant.1361, %constant.1361, %constant.1361), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.635 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1558 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.635), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.280 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1558, %convert_element_type.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1362 = f32[]{:T(128)} constant(0) - ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.280, %constant.1362), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} -} - -%fused_computation.179.clone.1.clone (param_0.1535: f32[4,128]) -> f32[4,128] { - %param_0.1535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1364 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1364), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1535, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1363 = f32[]{:T(128)} constant(1e-06) - %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1363), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} +%fused_computation.58.clone.clone (param_0.1535: bf16[4,4,128,2048], param_1.1688: s32[]) -> f32[4,128] { + %param_0.1535 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1688 = s32[]{:T(128)S(6)} parameter(1) + %constant.1363 = s32[]{:T(128)} constant(0) + %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1535, %param_1.1688, %constant.1363, %constant.1363, %constant.1363), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1564 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.633), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.280 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1564, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1364 = f32[]{:T(128)} constant(0) + ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.280, %constant.1364), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} +} + +%fused_computation.179.clone.1.clone (param_0.1536: f32[4,128]) -> f32[4,128] { + %param_0.1536 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %constant.1366 = f32[]{:T(128)} constant(0.00048828125) + %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1366), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1536, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1365 = f32[]{:T(128)} constant(1e-06) + %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1365), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.1039 = f32[4,128]{1,0:T(4,128)} add(%div.999, %closed_call.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.181 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1039), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } @@ -1548,158 +1548,158 @@ StackFrames ROOT %reduce_sum.212 = f32[]{:T(128)} add(%reduce_sum.207, %reduce_sum.211), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1549: bf16[4,2048,16,128], param_1.1691: s32[]) -> bf16[2048,16,128,1] { - %param_0.1549 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1691 = s32[]{:T(128)S(6)} parameter(1) - %constant.1375 = s32[]{:T(128)} constant(0) - %dynamic_slice.391 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1691, %constant.1375, %constant.1375, %constant.1375), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.646 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.114.clone.clone.clone.clone (param_0.1550: f32[4,128], param_1.1692: bf16[4,4,128,2048], param_2.1403: s32[], param_3.983: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.983 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.983), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1692 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1403 = s32[]{:T(128)S(6)} parameter(2) - %constant.1376 = s32[]{:T(128)} constant(0) - %dynamic_slice.392 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1692, %param_2.1403, %constant.1376, %constant.1376, %constant.1376), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.648 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1569 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.648), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1550 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2248 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2247 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1569, %mul.2248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1568 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2247), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.569 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.570, %convert_element_type.1568), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.647 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.61.clone.clone (param_0.1551: bf16[4,2048,16,128], param_1.1693: s32[], param_2.1404: f32[4,128], param_3.984: bf16[4,4,128,2048], param_4.603: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { - %param_2.1404 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.984 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1693 = s32[]{:T(128)S(6)} parameter(1) - %param_4.603 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1404, %param_3.984, %param_1.1693, %param_4.603), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1551 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1551, %param_1.1693), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.46.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1570 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.46.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.282 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1570, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1377 = f32[]{:T(128)} constant(0) - %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.282, %constant.1377), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.46.clone.3) -} - -%fused_computation.151.clone.1.clone (param_0.1552: f32[4,128,16]) -> f32[4,128,16] { - %param_0.1552 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1378 = f32[]{:T(128)} constant(0.0078125) - %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1378), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1552, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1379 = f32[]{:T(128)} constant(1e-06) - %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1379), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} +%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1550: bf16[4,2048,16,128], param_1.1698: s32[]) -> bf16[2048,16,128,1] { + %param_0.1550 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1698 = s32[]{:T(128)S(6)} parameter(1) + %constant.1377 = s32[]{:T(128)} constant(0) + %dynamic_slice.395 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1550, %param_1.1698, %constant.1377, %constant.1377, %constant.1377), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.644 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.114.clone.clone.clone.clone (param_0.1551: f32[4,128], param_1.1699: bf16[4,4,128,2048], param_2.1405: s32[], param_3.982: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.982 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.571 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1699 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1405 = s32[]{:T(128)S(6)} parameter(2) + %constant.1378 = s32[]{:T(128)} constant(0) + %dynamic_slice.396 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1699, %param_2.1405, %constant.1378, %constant.1378, %constant.1378), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.646 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.646), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2256 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1551), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2255 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %mul.2256), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2255), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.571, %convert_element_type.1574), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.645 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.61.clone.clone (param_0.1552: bf16[4,2048,16,128], param_1.1700: s32[], param_2.1406: f32[4,128], param_3.983: bf16[4,4,128,2048], param_4.604: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { + %param_2.1406 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.983 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1700 = s32[]{:T(128)S(6)} parameter(1) + %param_4.604 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1406, %param_3.983, %param_1.1700, %param_4.604), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1552 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1552, %param_1.1700), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.44.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %convert_element_type.1576 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.44.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.282 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1576, %convert_element_type.1576), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1379 = f32[]{:T(128)} constant(0) + %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.282, %constant.1379), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.44.clone.3) +} + +%fused_computation.151.clone.1.clone (param_0.1553: f32[4,128,16]) -> f32[4,128,16] { + %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) + %constant.1380 = f32[]{:T(128)} constant(0.0078125) + %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1380), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1553, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1381 = f32[]{:T(128)} constant(1e-06) + %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} %add.1043 = f32[4,128,16]{1,2,0:T(8,128)} add(%div.1001, %add.1044), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.183 = f32[4,128,16]{1,2,0:T(8,128)S(1)} rsqrt(%add.1043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.182.clone.clone (param_0.1548: bf16[4,128], param_1.1690: s32[]) -> bf16[128] { - %param_0.1548 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1690 = s32[]{:T(128)S(6)} parameter(1) - %constant.1374 = s32[]{:T(128)} constant(0) - %dynamic_slice.390 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1548, %param_1.1690, %constant.1374), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.645 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.182.clone.clone (param_0.1549: bf16[4,128], param_1.1697: s32[]) -> bf16[128] { + %param_0.1549 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) + %constant.1376 = s32[]{:T(128)} constant(0) + %dynamic_slice.394 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1697, %constant.1376), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.643 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.394), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.121.clone.1.clone (param_0.1553: f32[4,128,16], param_1.1694: bf16[4,128,16,128], param_2.1405: bf16[128]) -> bf16[4,128,16,128] { - %param_2.1405 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1405), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1694 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1572 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1694), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2250 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1553), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2249 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1572, %mul.2250), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1571 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.571 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.572, %convert_element_type.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.121.clone.1.clone (param_0.1554: f32[4,128,16], param_1.1701: bf16[4,128,16,128], param_2.1407: bf16[128]) -> bf16[4,128,16,128] { + %param_2.1407 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1407), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1701 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1578 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1554 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) + %mul.2258 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1554), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2257 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1578, %mul.2258), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1577 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2257), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.573, %convert_element_type.1577), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} } -%fused_computation.90.clone.clone (param_0.1554: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { - %param_0.1554 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1554), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.90.clone.clone (param_0.1555: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { + %param_0.1555 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.129 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1554), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.209 = (bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) } %fused_computation.187.clone.clone () -> f32[64] { - %constant.1353 = f32[]{:T(128)} constant(1e+06) - %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1353), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %constant.1355 = f32[]{:T(128)} constant(1e+06) + %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1355), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} - %constant.1352 = s32[]{:T(128)} constant(2) - %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1352), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2234 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1556 = f32[64]{0:T(128)} convert(%mul.2234), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %constant.1354 = f32[]{:T(128)} constant(0.0078125) - %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1556, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1354 = s32[]{:T(128)} constant(2) + %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %mul.2242 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1562 = f32[64]{0:T(128)} convert(%mul.2242), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %constant.1356 = f32[]{:T(128)} constant(0.0078125) + %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1356), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1562, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.104, %div.995), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} } -%fused_computation.143.clone.clone (param_0.1528: f32[64], param_1.1676: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1676 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1676), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1528 = f32[64]{0:T(128)S(1)} parameter(0) - %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1528), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} +%fused_computation.143.clone.clone (param_0.1529: f32[64], param_1.1683: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { + %param_1.1683 = f32[4,128]{1,0:T(4,128)} parameter(1) + %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1683), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %param_0.1529 = f32[64]{0:T(128)S(1)} parameter(0) + %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1529), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %div.996 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.998, %div.997), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1557 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convert_element_type.1563 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} %convert_element_type.1189.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1557, %convert_element_type.1189.clone.3) + ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1563, %convert_element_type.1189.clone.3) } -%fused_computation.146.clone.1.clone (param_0.1529: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1529 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1355 = bf16[]{:T(256)} constant(-inf) - %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1529, %constant.1355), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1529, %constant.1355), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} +%fused_computation.146.clone.1.clone (param_0.1530: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1530 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) + %constant.1357 = bf16[]{:T(256)} constant(-inf) + %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.53 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.69, %pad.68), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.632 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.630 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.145.clone.1.clone (param_0.1544: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1544 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1372 = bf16[]{:T(256)} constant(-inf) - %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1544, %constant.1372), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1544, %constant.1372), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} +%fused_computation.145.clone.1.clone (param_0.1545: bf16[4,128,1,64]) -> bf16[4,128,128] { + %param_0.1545 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) + %constant.1374 = bf16[]{:T(256)} constant(-inf) + %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.54 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.71, %pad.70), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.643 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.641 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} } -%fused_computation.94.clone.clone (param_0.1555: bf16[4,128,16,64], param_1.1695: bf16[4,128,16,64], param_2.1406: bf16[4,128,128], param_3.985: bf16[4,128,128], param_4.604: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.383: bf16[128]) -> bf16[4,16,128,128] { - %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.94.clone.clone (param_0.1556: bf16[4,128,16,64], param_1.1702: bf16[4,128,16,64], param_2.1408: bf16[4,128,128], param_3.984: bf16[4,128,128], param_4.605: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.384: bf16[128]) -> bf16[4,16,128,128] { + %param_6.384 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.575 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.384), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} %param_5.499 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1574 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.604 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2257 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.604), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2256 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1574, %mul.2257), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2256), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %convert_element_type.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.985 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2255 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.985), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2253 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.573, %mul.2255), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1695 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1380 = bf16[]{:T(256)} constant(-inf) - %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1695, %constant.1380), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1555 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1555, %constant.1380), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %convert_element_type.1580 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_4.605 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) + %mul.2265 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.605), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2264 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1580, %mul.2265), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1579 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.575, %convert_element_type.1579), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.984 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2263 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.984), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2261 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %mul.2263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1702 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1382 = bf16[]{:T(256)} constant(-inf) + %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1702, %constant.1382), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1556 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1556, %constant.1382), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.56 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.75, %pad.74), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1406 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2254 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1406), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2252 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2254), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2253, %mul.2252), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %constant.1381 = bf16[]{:T(256)} constant(0.08838) - %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2251 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - ROOT %bitcast.649 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_2.1408 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2262 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1408), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2260 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2262), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2261, %mul.2260), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + %constant.1383 = bf16[]{:T(256)} constant(0.08838) + %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1383), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %mul.2259 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + ROOT %bitcast.647 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2259), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } %region_16.18 (reduce_sum.213: f32[], reduce_sum.214: f32[]) -> f32[] { @@ -1708,159 +1708,159 @@ StackFrames ROOT %reduce_sum.218 = f32[]{:T(128)} add(%reduce_sum.213, %reduce_sum.214), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1540: bf16[4,2048,8,128], param_1.1685: s32[]) -> bf16[2048,8,128,1] { - %param_0.1540 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) - %constant.1367 = s32[]{:T(128)} constant(0) - %dynamic_slice.388 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1540, %param_1.1685, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.640 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.113.clone.clone.clone.clone (param_0.1541: f32[4,128], param_1.1686: bf16[4,4,128,2048], param_2.1399: s32[], param_3.980: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.980 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.980), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1686 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) - %constant.1368 = s32[]{:T(128)} constant(0) - %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1686, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.642 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1562 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.642), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1541 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2238 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1541), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2237 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1562, %mul.2238), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1561 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2237), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.564, %convert_element_type.1561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.641 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.84.clone.clone (param_0.1542: bf16[4,2048,8,128], param_1.1687: s32[], param_2.1400: f32[4,128], param_3.981: bf16[4,4,128,2048], param_4.601: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { - %param_2.1400 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.981 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) - %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1400, %param_3.981, %param_1.1687, %param_4.601), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1542 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1542, %param_1.1687), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1541: bf16[4,2048,8,128], param_1.1692: s32[]) -> bf16[2048,8,128,1] { + %param_0.1541 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1692 = s32[]{:T(128)S(6)} parameter(1) + %constant.1369 = s32[]{:T(128)} constant(0) + %dynamic_slice.392 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1541, %param_1.1692, %constant.1369, %constant.1369, %constant.1369), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.638 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.113.clone.clone.clone.clone (param_0.1542: f32[4,128], param_1.1693: bf16[4,4,128,2048], param_2.1401: s32[], param_3.979: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.979 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.979), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1693 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1401 = s32[]{:T(128)S(6)} parameter(2) + %constant.1370 = s32[]{:T(128)} constant(0) + %dynamic_slice.393 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1693, %param_2.1401, %constant.1370, %constant.1370, %constant.1370), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.640 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1568 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.640), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1542 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2246 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1542), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2245 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1568, %mul.2246), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.565, %convert_element_type.1567), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.639 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.84.clone.clone (param_0.1543: bf16[4,2048,8,128], param_1.1694: s32[], param_2.1402: f32[4,128], param_3.980: bf16[4,4,128,2048], param_4.602: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { + %param_2.1402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.980 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1694 = s32[]{:T(128)S(6)} parameter(1) + %param_4.602 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1402, %param_3.980, %param_1.1694, %param_4.602), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1543 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1543, %param_1.1694), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %convolution.50.clone.3 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.73.clone.3, %fusion.87.clone.3), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1563 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.281 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1563, %convert_element_type.1563), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1369 = f32[]{:T(128)} constant(0) - %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.281, %constant.1369), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %convert_element_type.1569 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.281 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1569, %convert_element_type.1569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1371 = f32[]{:T(128)} constant(0) + %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.281, %constant.1371), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.206 = (f32[4,128,8]{1,2,0:T(8,128)S(1)}, bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.246, %convolution.50.clone.3) } -%fused_computation.154.clone.1.clone (param_0.1543: f32[4,128,8]) -> f32[4,128,8] { - %param_0.1543 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1370 = f32[]{:T(128)} constant(0.0078125) - %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1370), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1543, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1371 = f32[]{:T(128)} constant(1e-06) - %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1371), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} +%fused_computation.154.clone.1.clone (param_0.1544: f32[4,128,8]) -> f32[4,128,8] { + %param_0.1544 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) + %constant.1372 = f32[]{:T(128)} constant(0.0078125) + %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1372), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1544, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1373 = f32[]{:T(128)} constant(1e-06) + %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1373), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} %add.1040 = f32[4,128,8]{1,2,0:T(8,128)} add(%div.1000, %add.1041), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.182 = f32[4,128,8]{1,2,0:T(8,128)S(1)} rsqrt(%add.1040), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.184.clone.clone (param_0.1527: bf16[4,128], param_1.1675: s32[]) -> bf16[128] { - %param_0.1527 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1675 = s32[]{:T(128)S(6)} parameter(1) - %constant.1351 = s32[]{:T(128)} constant(0) - %dynamic_slice.381 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1527, %param_1.1675, %constant.1351), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.631 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.139.clone.1.clone (param_0.1545: f32[4,128,8], param_1.1688: bf16[4,128,8,128], param_2.1401: bf16[128]) -> bf16[4,128,8,128] { - %param_2.1401 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1401), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1688 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1565 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1688), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1545 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2240 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1545), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2239 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1565, %mul.2240), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1564 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2239), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.565 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.566, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.126.clone.clone (param_0.1546: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1546 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1546), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} +%fused_computation.184.clone.clone (param_0.1528: bf16[4,128], param_1.1682: s32[]) -> bf16[128] { + %param_0.1528 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) + %constant.1353 = s32[]{:T(128)} constant(0) + %dynamic_slice.385 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1528, %param_1.1682, %constant.1353), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.629 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.139.clone.1.clone (param_0.1546: f32[4,128,8], param_1.1695: bf16[4,128,8,128], param_2.1403: bf16[128]) -> bf16[4,128,8,128] { + %param_2.1403 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1403), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1695 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1571 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1695), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1546 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) + %mul.2248 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1546), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2247 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1571, %mul.2248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1570 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2247), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.567, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.126.clone.clone (param_0.1547: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { + %param_0.1547 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1546), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} + %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} ROOT %tuple.207 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) } -%fused_computation.129.clone.clone (param_0.1547: bf16[4,128,8,64], param_1.1689: bf16[4,128,8,64], param_2.1402: bf16[4,128,128], param_3.982: bf16[4,128,128], param_4.602: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.382: bf16[128]) -> bf16[4,8,128,128] { - %param_6.382 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.382), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.129.clone.clone (param_0.1548: bf16[4,128,8,64], param_1.1696: bf16[4,128,8,64], param_2.1404: bf16[4,128,128], param_3.981: bf16[4,128,128], param_4.603: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.383: bf16[128]) -> bf16[4,8,128,128] { + %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.569 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} %param_5.498 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1567 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.602 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2246 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.602), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2245 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1567, %mul.2246), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %convert_element_type.1566), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.982 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2244 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2242 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.567, %mul.2244), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1689 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1373 = bf16[]{:T(256)} constant(-inf) - %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1689, %constant.1373), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1547 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1547, %constant.1373), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %convert_element_type.1573 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_4.603 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) + %mul.2254 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.603), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2253 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1573, %mul.2254), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1572 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2253), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.569, %convert_element_type.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.981 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2252 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.981), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2250 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %mul.2252), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %param_1.1696 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1375 = bf16[]{:T(256)} constant(-inf) + %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1696, %constant.1375), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} + %param_0.1548 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1548, %constant.1375), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} %maximum.55 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.73, %pad.72), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1402 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2243 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1402), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2241 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2243), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2242, %mul.2241), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.644 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_2.1404 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2251 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1404), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2249 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2250, %mul.2249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} + ROOT %bitcast.642 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.169.clone.clone (param_0.1536: bf16[4,2048,8,128], param_1.1682: s32[]) -> bf16[1,2048,8,128] { - %param_0.1536 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) - %constant.1365 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.386 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1536, %param_1.1682, %constant.1365, %constant.1365, %constant.1365), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} -} - -%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1537: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { - %param_0.1537 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.208 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1537), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.636 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.208), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.111.clone.clone.clone.clone (param_0.1538: f32[4,128], param_1.1683: bf16[4,4,128,2048], param_2.1397: s32[], param_3.978: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.978 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.978), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1683 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1397 = s32[]{:T(128)S(6)} parameter(2) - %constant.1366 = s32[]{:T(128)} constant(0) - %dynamic_slice.387 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1683, %param_2.1397, %constant.1366, %constant.1366, %constant.1366), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.638 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1560 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.638), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1538 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2236 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1538), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2235 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1560, %mul.2236), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1559 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2235), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.561 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.562, %convert_element_type.1559), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.637 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.140.clone.clone (param_0.1539: bf16[1,2048,8,128], param_1.1684: f32[4,128], param_2.1398: bf16[4,4,128,2048], param_3.979: s32[], param_4.600: bf16[2048]) -> bf16[4,8,128,128] { - %param_1.1684 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1398 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.979 = s32[]{:T(128)S(6)} parameter(3) - %param_4.600 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.372 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1684, %param_2.1398, %param_3.979, %param_4.600), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1539 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.371 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1539), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.372, %fusion.371), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.639 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.188.clone.clone (param_0.1577: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { - %param_0.1577 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1577), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.662 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.662, %slice.11) +%fused_computation.169.clone.clone (param_0.1537: bf16[4,2048,8,128], param_1.1689: s32[]) -> bf16[1,2048,8,128] { + %param_0.1537 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) + %param_1.1689 = s32[]{:T(128)S(6)} parameter(1) + %constant.1367 = s32[]{:T(128)} constant(0) + ROOT %dynamic_slice.390 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1537, %param_1.1689, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1538: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { + %param_0.1538 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %copy.204 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1538), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} + ROOT %bitcast.634 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.204), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.111.clone.clone.clone.clone (param_0.1539: f32[4,128], param_1.1690: bf16[4,4,128,2048], param_2.1399: s32[], param_3.977: bf16[2048]) -> bf16[4,128,2048,1] { + %param_3.977 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.977), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1690 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) + %constant.1368 = s32[]{:T(128)} constant(0) + %dynamic_slice.391 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1690, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.636 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1566 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.636), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1539 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2244 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2243 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1566, %mul.2244), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2243), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.563, %convert_element_type.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %bitcast.635 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.562), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.140.clone.clone (param_0.1540: bf16[1,2048,8,128], param_1.1691: f32[4,128], param_2.1400: bf16[4,4,128,2048], param_3.978: s32[], param_4.601: bf16[2048]) -> bf16[4,8,128,128] { + %param_1.1691 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) + %param_2.1400 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.978 = s32[]{:T(128)S(6)} parameter(3) + %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.373 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1691, %param_2.1400, %param_3.978, %param_4.601), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1540 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.372 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1540), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.373, %fusion.372), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + ROOT %bitcast.637 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +} + +%fused_computation.188.clone.clone (param_0.1578: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { + %param_0.1578 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) + %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1578), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.660 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} + ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.660, %slice.11) } %region_17.20 (reduce_sum.219: f32[], reduce_sum.220: f32[]) -> f32[] { @@ -1869,36 +1869,36 @@ StackFrames ROOT %reduce_sum.221 = f32[]{:T(128)} add(%reduce_sum.219, %reduce_sum.220), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1556: bf16[4,16,128,2048], param_1.1696: s32[]) -> bf16[16,128,2048,1] { - %param_0.1556 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1696 = s32[]{:T(128)S(6)} parameter(1) - %constant.1382 = s32[]{:T(128)} constant(0) - %dynamic_slice.393 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1556, %param_1.1696, %constant.1382, %constant.1382, %constant.1382), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.650 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,2048], param_1.1703: s32[]) -> bf16[16,128,2048,1] { + %param_0.1557 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) + %constant.1384 = s32[]{:T(128)} constant(0) + %dynamic_slice.397 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1557, %param_1.1703, %constant.1384, %constant.1384, %constant.1384), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.648 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.397), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,128]) -> bf16[4,128,16,128] { - %param_0.1557 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.651 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} +%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1558: bf16[4,16,128,128]) -> bf16[4,128,16,128] { + %param_0.1558 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.649 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} } -%fused_computation.64.clone.clone (param_0.1558: bf16[4,16,128,2048], param_1.1697: s32[], param_2.1407: bf16[4,16,128,128], param_3.986: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { - %param_3.986 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.64.clone.clone (param_0.1559: bf16[4,16,128,2048], param_1.1704: s32[], param_2.1409: bf16[4,16,128,128], param_3.985: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { + %param_3.985 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) + %param_1.1704 = s32[]{:T(128)S(6)} parameter(1) %constant.436.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.240.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.986, %param_1.1697, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.240.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1407 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1407), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1558 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1558, %param_1.1697), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %dynamic_slice.242.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.985, %param_1.1704, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.242.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} + %param_2.1409 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1409), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %param_0.1559 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1559, %param_1.1704), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} %convolution.62.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.96.clone.3, %fusion.95.clone.3), window={size=1x16}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %bitcast.203.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} %add.768.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.227.clone.3, %bitcast.203.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.283 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %convert_element_type.1575), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1383 = f32[]{:T(128)} constant(0) - %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.283, %constant.1383), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} + %convert_element_type.1581 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %square.283 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1581, %convert_element_type.1581), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} + %constant.1385 = f32[]{:T(128)} constant(0) + %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.283, %constant.1385), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} ROOT %tuple.210 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.248, %add.768.clone.3) } @@ -1908,93 +1908,93 @@ StackFrames ROOT %add.754 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.156.clone.clone (param_0.1530: bf16[4,2048], param_1.1677: s32[]) -> bf16[2048] { - %param_0.1530 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1677 = s32[]{:T(128)S(6)} parameter(1) - %constant.1356 = s32[]{:T(128)} constant(0) - %dynamic_slice.382 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1530, %param_1.1677, %constant.1356), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1357 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.382, %constant.1357), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.156.clone.clone (param_0.1531: bf16[4,2048], param_1.1684: s32[]) -> bf16[2048] { + %param_0.1531 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1684 = s32[]{:T(128)S(6)} parameter(1) + %constant.1358 = s32[]{:T(128)} constant(0) + %dynamic_slice.386 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1684, %constant.1358), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1359 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.386, %constant.1359), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } -%fused_computation.13.clone.clone.clone (param_0.1531: bf16[4,6144,2048], param_1.1678: s32[]) -> bf16[6144,2048,1] { - %param_0.1531 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1678 = s32[]{:T(128)S(6)} parameter(1) - %constant.1358 = s32[]{:T(128)} constant(0) - %dynamic_slice.383 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1678, %constant.1358, %constant.1358), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.634 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.13.clone.clone.clone (param_0.1532: bf16[4,6144,2048], param_1.1685: s32[]) -> bf16[6144,2048,1] { + %param_0.1532 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) + %constant.1360 = s32[]{:T(128)} constant(0) + %dynamic_slice.387 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1532, %param_1.1685, %constant.1360, %constant.1360), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.632 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} } %bitcast_fusion.1.clone.clone (bitcast_input.4: bf16[4,128,2048]) -> bf16[4,128,2048] { - %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) -} - -%fused_computation.14.clone.clone (param_0.1532: bf16[4,128,2048], param_1.1679: bf16[4,6144,2048], param_2.1396: s32[]) -> bf16[6144,4,128] { - %param_1.1679 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1396 = s32[]{:T(128)S(6)} parameter(2) - %fusion.369 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1679, %param_2.1396), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1532 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.370 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1532), kind=kLoop, calls=%bitcast_fusion.1.clone.clone - ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.369, %fusion.370), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.180.clone.1.clone (param_0.1559: f32[4,128]) -> f32[4,128] { - %param_0.1559 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1385 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1385), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1559, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1384 = f32[]{:T(128)} constant(1e-06) - %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1384), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.631 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) +} + +%fused_computation.14.clone.clone (param_0.1533: bf16[4,128,2048], param_1.1686: bf16[4,6144,2048], param_2.1398: s32[]) -> bf16[6144,4,128] { + %param_1.1686 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) + %param_2.1398 = s32[]{:T(128)S(6)} parameter(2) + %fusion.370 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1686, %param_2.1398), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1533 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.371 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1533), kind=kLoop, calls=%bitcast_fusion.1.clone.clone + ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.370, %fusion.371), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.180.clone.1.clone (param_0.1560: f32[4,128]) -> f32[4,128] { + %param_0.1560 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %constant.1387 = f32[]{:T(128)} constant(0.00048828125) + %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1387), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} + %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1560, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} + %constant.1386 = f32[]{:T(128)} constant(1e-06) + %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1386), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} %add.1046 = f32[4,128]{1,0:T(4,128)} add(%div.1002, %closed_call.110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} ROOT %rsqrt.184 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1046), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} } -%fused_computation.12.clone.1.clone.clone (param_0.1563: bf16[4,2048,6144], param_1.1701: s32[]) -> bf16[2048,6144,1] { - %param_0.1563 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1701 = s32[]{:T(128)S(6)} parameter(1) - %constant.1387 = s32[]{:T(128)} constant(0) - %dynamic_slice.395 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1563, %param_1.1701, %constant.1387, %constant.1387), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.119.clone.3.clone.clone (param_0.1564: f32[4,128], param_1.1702: bf16[4,128,2048], param_2.1410: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1410 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1410), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1702 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1579 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1564 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2261 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1564), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2260 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1579, %mul.2261), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2260), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.577 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.578, %convert_element_type.1578), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.21.clone.clone (param_0.1565: bf16[4,2048,6144], param_1.1703: s32[], param_2.1411: f32[4,128], param_3.988: bf16[4,128,2048], param_4.606: bf16[2048]) -> bf16[4,128,6144] { - %param_2.1411 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.988 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.606 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.376 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1411, %param_3.988, %param_4.606), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1565 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) - %fusion.375 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1565, %param_1.1703), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.376, %fusion.375), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.11.clone.1.clone.clone (param_0.1567: bf16[4,2048,6144], param_1.1705: s32[]) -> bf16[2048,6144,1] { - %param_0.1567 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1705 = s32[]{:T(128)S(6)} parameter(1) +%fused_computation.12.clone.1.clone.clone (param_0.1564: bf16[4,2048,6144], param_1.1708: s32[]) -> bf16[2048,6144,1] { + %param_0.1564 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1708 = s32[]{:T(128)S(6)} parameter(1) %constant.1389 = s32[]{:T(128)} constant(0) - %dynamic_slice.396 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1567, %param_1.1705, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.655 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.47.clone.1.clone.clone (param_0.1566: bf16[6144,4,128], param_1.1704: bf16[4,128,6144]) -> bf16[4,128,6144] { - %param_1.1704 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1388 = bf16[]{:T(256)} constant(1) - %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1388), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} + %dynamic_slice.399 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1564, %param_1.1708, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.651 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.119.clone.3.clone.clone (param_0.1565: f32[4,128], param_1.1709: bf16[4,128,2048], param_2.1412: bf16[2048]) -> bf16[4,128,2048] { + %param_2.1412 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.579 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1412), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1709 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1585 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + %param_0.1565 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) + %mul.2269 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1565), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %mul.2268 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1585, %mul.2269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} + %convert_element_type.1584 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.579, %convert_element_type.1584), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.21.clone.clone (param_0.1566: bf16[4,2048,6144], param_1.1710: s32[], param_2.1413: f32[4,128], param_3.987: bf16[4,128,2048], param_4.607: bf16[2048]) -> bf16[4,128,6144] { + %param_2.1413 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) + %param_3.987 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.607 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.377 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1413, %param_3.987, %param_4.607), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} + %param_0.1566 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1710 = s32[]{:T(128)S(6)} parameter(1) + %fusion.376 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1566, %param_1.1710), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} + ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.377, %fusion.376), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} +} + +%fused_computation.11.clone.1.clone.clone (param_0.1568: bf16[4,2048,6144], param_1.1712: s32[]) -> bf16[2048,6144,1] { + %param_0.1568 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1712 = s32[]{:T(128)S(6)} parameter(1) + %constant.1391 = s32[]{:T(128)} constant(0) + %dynamic_slice.400 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1568, %param_1.1712, %constant.1391, %constant.1391), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.400), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +} + +%fused_computation.47.clone.1.clone.clone (param_0.1567: bf16[6144,4,128], param_1.1711: bf16[4,128,6144]) -> bf16[4,128,6144] { + %param_1.1711 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1390 = bf16[]{:T(256)} constant(1) + %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1390), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} + %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1711), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} %exp.69 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} %add.1047 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} %div.1003 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.1047), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.2263 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1704, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %mul.2271 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1711, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} From 15305a710caaa02cdab68d191cb4bf4af3bf6250 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 11 Jun 2026 16:07:03 +0000 Subject: [PATCH 33/52] Fix pytest marker for HLO diff test to use tpu_only --- tests/integration/hlo_diff_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/hlo_diff_test.py b/tests/integration/hlo_diff_test.py index cc18713749..8077a9e83f 100644 --- a/tests/integration/hlo_diff_test.py +++ b/tests/integration/hlo_diff_test.py @@ -59,7 +59,7 @@ def filter_line(line): not os.environ.get("GITHUB_ACTIONS"), reason="Skipping HLO diff test because it is not running in GitHub Actions", ) -@pytest.mark.tpu_backend +@pytest.mark.tpu_only class TestHloDiff: """Tests for HLO Graph Diff Verification.""" From 739cb6602b6beb98695debc59d10e68b11cd2842 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 11 Jun 2026 17:31:52 +0000 Subject: [PATCH 34/52] test: mark data loader and grain tests as cpu_only to fix multi-device deadlocks --- tests/unit/data_loader_test.py | 2 ++ tests/unit/grain_data_processing_test.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/tests/unit/data_loader_test.py b/tests/unit/data_loader_test.py index c61217d598..16c1b2e193 100644 --- a/tests/unit/data_loader_test.py +++ b/tests/unit/data_loader_test.py @@ -26,6 +26,8 @@ from maxtext.configs import pyconfig from maxtext.common.data_loader import DataLoader, RampUpDataLoader + +pytestmark = pytest.mark.cpu_only from maxtext.utils import exceptions from maxtext.utils.maxtext_utils import create_device_mesh from maxtext.common.gcloud_stub import is_decoupled diff --git a/tests/unit/grain_data_processing_test.py b/tests/unit/grain_data_processing_test.py index dc129b5ebc..30cff4e46c 100644 --- a/tests/unit/grain_data_processing_test.py +++ b/tests/unit/grain_data_processing_test.py @@ -29,6 +29,8 @@ from maxtext.configs import pyconfig from maxtext.input_pipeline import grain_data_processing from maxtext.input_pipeline import input_pipeline_interface + +pytestmark = pytest.mark.cpu_only from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from maxtext.common.gcloud_stub import is_decoupled from tests.utils.test_helpers import get_test_base_output_directory, get_test_config_path, get_test_dataset_path From 74cf87fc0e848c3bdb0599687b51da4de79f0cf9 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 11 Jun 2026 21:21:20 +0000 Subject: [PATCH 35/52] Fix NCCL invalid argument failures in multi-GPU CI tests Removed NCCL_SOCKET_IFNAME=lo and NCCL_NET_GDR_LEVEL=0 to allow NCCL to auto-discover the optimal network interface and transport (NVLink/P2P) in the Docker container, fixing ncclInvalidArgument errors. Also changed NCCL_DEBUG to WARN to prevent log spam. TAG=agy CONV=bf394eb7-23a3-4437-91ef-0cead1a5b0a0 --- .github/workflows/run_tests_against_package.yml | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 491beafe8b..ccc5d01697 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -181,11 +181,10 @@ jobs: echo "Final LD_LIBRARY_PATH: ${LD_LIBRARY_PATH}" echo "=======================" - # Configure NCCL for stable single-node communication in Docker containers - export NCCL_SOCKET_IFNAME=lo - export NCCL_NET_GDR_LEVEL=0 - export NCCL_DEBUG=INFO - echo "Set NCCL_SOCKET_IFNAME=lo, NCCL_NET_GDR_LEVEL=0, and NCCL_DEBUG=INFO for GPU execution." + # Configure NCCL for GPU execution + # Removed NCCL_SOCKET_IFNAME=lo as it breaks multi-gpu collective ops in Docker + export NCCL_DEBUG=WARN + echo "Set NCCL_DEBUG=WARN for GPU execution." fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist From fdf4ed976ee1c810ea18aaa3ebde806c2321a876 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 11 Jun 2026 21:52:44 +0000 Subject: [PATCH 36/52] Remove LD_LIBRARY_PATH hack to fix NCCL comm corruption --- .../workflows/run_tests_against_package.yml | 22 +------------------ 1 file changed, 1 insertion(+), 21 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index ccc5d01697..8ab8b4e078 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -159,27 +159,7 @@ jobs: if [ "${INPUTS_DEVICE_TYPE}" != "cuda12" ]; then export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' else - # For cuda12, explicitly point to the pip-installed CUDA libraries - # to avoid conflicts with system-level installations on the runner. - # Dynamically discover the 'nvidia' folder and prepend all its sub-library - # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent - # JAX from partially loading incompatible system-level CUDA libraries. - NVIDIA_DIR=$(find -L .venv -path "*/site-packages/nvidia" -type d 2>/dev/null | head -n 1) - echo "=== GPU Diagnostics ===" - echo "Found NVIDIA_DIR: ${NVIDIA_DIR}" - if [ -n "${NVIDIA_DIR}" ]; then - for dir in "${NVIDIA_DIR}"/*; do - if [ -d "$dir/lib" ]; then - ABS_LIB_PATH=$(realpath "$dir/lib") - export LD_LIBRARY_PATH=${ABS_LIB_PATH}:${LD_LIBRARY_PATH} - echo "Prepended to LD_LIBRARY_PATH: ${ABS_LIB_PATH}" - fi - done - else - echo "WARNING: nvidia directory not found under .venv!" - fi - echo "Final LD_LIBRARY_PATH: ${LD_LIBRARY_PATH}" - echo "=======================" + # Configure NCCL for GPU execution # Removed NCCL_SOCKET_IFNAME=lo as it breaks multi-gpu collective ops in Docker From 197c41fab7a0326f10cb25551170ba44d6c4d625 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 13 Jun 2026 04:56:06 +0000 Subject: [PATCH 37/52] Fix GPU runner CI failures: restore LD_LIBRARY_PATH, exclude tests, catch OSError --- .github/workflows/run_tests_against_package.yml | 14 +++++++++++++- tests/unit/data_loader_test.py | 3 ++- tests/unit/grain_data_processing_test.py | 3 ++- tests/unit/hf_data_processing_test.py | 3 +++ tests/unit/quantizations_test.py | 2 +- tests/unit/tfds_data_processing_test.py | 2 ++ 6 files changed, 23 insertions(+), 4 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 8ab8b4e078..9bccc54e7a 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -159,7 +159,19 @@ jobs: if [ "${INPUTS_DEVICE_TYPE}" != "cuda12" ]; then export LIBTPU_INIT_ARGS='--xla_tpu_scoped_vmem_limit_kib=65536' else - + # For cuda12, explicitly point to the pip-installed CUDA libraries + # to avoid conflicts with system-level installations on the runner. + # Dynamically discover the 'nvidia' folder and prepend all its sub-library + # directories (including nccl, cublas, cudnn) to LD_LIBRARY_PATH to prevent + # JAX from partially loading incompatible system-level CUDA libraries. + NVIDIA_DIR=$(find .venv/lib/ -maxdepth 3 -name "nvidia" -type d 2>/dev/null | head -n 1) + if [ -n "${NVIDIA_DIR}" ]; then + for dir in "${NVIDIA_DIR}"/*; do + if [ -d "$dir/lib" ]; then + export LD_LIBRARY_PATH=$(pwd)/$dir/lib:${LD_LIBRARY_PATH} + fi + done + fi # Configure NCCL for GPU execution # Removed NCCL_SOCKET_IFNAME=lo as it breaks multi-gpu collective ops in Docker diff --git a/tests/unit/data_loader_test.py b/tests/unit/data_loader_test.py index 16c1b2e193..52737c11c4 100644 --- a/tests/unit/data_loader_test.py +++ b/tests/unit/data_loader_test.py @@ -27,13 +27,14 @@ from maxtext.configs import pyconfig from maxtext.common.data_loader import DataLoader, RampUpDataLoader -pytestmark = pytest.mark.cpu_only from maxtext.utils import exceptions from maxtext.utils.maxtext_utils import create_device_mesh from maxtext.common.gcloud_stub import is_decoupled from maxtext.utils.rampup_batch import RampupBatchManager from tests.utils.test_helpers import get_test_config_path +pytestmark = pytest.mark.cpu_only + class DataLoaderTest(unittest.TestCase): diff --git a/tests/unit/grain_data_processing_test.py b/tests/unit/grain_data_processing_test.py index 30cff4e46c..1c22af1eac 100644 --- a/tests/unit/grain_data_processing_test.py +++ b/tests/unit/grain_data_processing_test.py @@ -30,11 +30,12 @@ from maxtext.input_pipeline import grain_data_processing from maxtext.input_pipeline import input_pipeline_interface -pytestmark = pytest.mark.cpu_only from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from maxtext.common.gcloud_stub import is_decoupled from tests.utils.test_helpers import get_test_base_output_directory, get_test_config_path, get_test_dataset_path +pytestmark = pytest.mark.cpu_only + class GrainBaseProcessingTest: """Base mixin with test_train_ds for all grain data processing tests. diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 262c56ff9b..224e803804 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -29,6 +29,9 @@ from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory +import pytest +pytestmark = pytest.mark.cpu_only + class HfDataProcessingTest(unittest.TestCase): diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 037c9e6f2c..d50ad9b6b8 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -356,7 +356,7 @@ def test_configure_quantization_paths(self): import transformer_engine # pylint: disable=unused-import,import-outside-toplevel has_te = True - except ImportError: + except (ImportError, OSError): has_te = False if has_te: diff --git a/tests/unit/tfds_data_processing_test.py b/tests/unit/tfds_data_processing_test.py index fe1bb75f24..1952391868 100644 --- a/tests/unit/tfds_data_processing_test.py +++ b/tests/unit/tfds_data_processing_test.py @@ -32,6 +32,8 @@ from maxtext.input_pipeline import input_pipeline_interface from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory +pytestmark = pytest.mark.cpu_only + class TfdsDataProcessingTest(unittest.TestCase): From 82e278cef419f528ebc324d56e47d6dccb14aa4c Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 13 Jun 2026 05:02:45 +0000 Subject: [PATCH 38/52] Fix pyink formatting in hf_data_processing_test.py --- tests/unit/hf_data_processing_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 224e803804..2a6cf93cd2 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -30,6 +30,7 @@ from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory import pytest + pytestmark = pytest.mark.cpu_only From 829283b338b69f1dfd55f8f8c5b535edd99e9fbd Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 13 Jun 2026 05:38:39 +0000 Subject: [PATCH 39/52] Temporarily enable triggering tests on push for deprecate-aqt-keep-deps --- .github/workflows/CodeQuality.yml | 3 +++ .github/workflows/build_and_test_maxtext.yml | 3 +++ 2 files changed, 6 insertions(+) diff --git a/.github/workflows/CodeQuality.yml b/.github/workflows/CodeQuality.yml index 9b62b3d507..7cb2f5d925 100644 --- a/.github/workflows/CodeQuality.yml +++ b/.github/workflows/CodeQuality.yml @@ -15,6 +15,9 @@ name: CodeQuality on: + push: + branches: + - deprecate-aqt-keep-deps pull_request: concurrency: diff --git a/.github/workflows/build_and_test_maxtext.yml b/.github/workflows/build_and_test_maxtext.yml index d1e5248f52..f65554416d 100644 --- a/.github/workflows/build_and_test_maxtext.yml +++ b/.github/workflows/build_and_test_maxtext.yml @@ -17,6 +17,9 @@ name: MaxText Package Tests on: + push: + branches: + - deprecate-aqt-keep-deps pull_request: workflow_call: workflow_dispatch: From bb95b7c153ed9e1e0e52a9316aa6d6712d9edfaa Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sat, 13 Jun 2026 05:45:11 +0000 Subject: [PATCH 40/52] Fix CodeQuality workflow for push events --- .github/workflows/CodeQuality.yml | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/.github/workflows/CodeQuality.yml b/.github/workflows/CodeQuality.yml index 7cb2f5d925..159ea8ad1d 100644 --- a/.github/workflows/CodeQuality.yml +++ b/.github/workflows/CodeQuality.yml @@ -57,7 +57,9 @@ jobs: - name: Run pre-commit checks on just the files that have changed run: | - git fetch origin "$GITHUB_BASE_REF":"$GITHUB_BASE_REF" - git branch "$GITHUB_HEAD_REF" + BASE_REF="${GITHUB_BASE_REF:-main}" + HEAD_REF="${GITHUB_HEAD_REF:-$GITHUB_SHA}" + git fetch origin "$BASE_REF":"$BASE_REF" + git branch "$HEAD_REF" || true . "$GITHUB_WORKSPACE"/venv/bin/activate - pre-commit run --from-ref "$GITHUB_BASE_REF" --to-ref "$GITHUB_HEAD_REF" --show-diff-on-failure + pre-commit run --from-ref "$BASE_REF" --to-ref "$HEAD_REF" --show-diff-on-failure From 70f534f924c1eb496cd702b73efb598f0922e5b0 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sun, 14 Jun 2026 00:50:17 +0000 Subject: [PATCH 41/52] Fix GPU CI: Remove duplicate workflow push triggers and add cpu_only to grain forking tests --- .github/workflows/build_and_test_maxtext.yml | 3 --- tests/unit/input_pipeline/olmo_data_grain_resume_test.py | 4 ++++ tests/unit/input_pipeline/olmo_data_grain_test.py | 4 ++++ 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build_and_test_maxtext.yml b/.github/workflows/build_and_test_maxtext.yml index f65554416d..d1e5248f52 100644 --- a/.github/workflows/build_and_test_maxtext.yml +++ b/.github/workflows/build_and_test_maxtext.yml @@ -17,9 +17,6 @@ name: MaxText Package Tests on: - push: - branches: - - deprecate-aqt-keep-deps pull_request: workflow_call: workflow_dispatch: diff --git a/tests/unit/input_pipeline/olmo_data_grain_resume_test.py b/tests/unit/input_pipeline/olmo_data_grain_resume_test.py index bc1f547496..460ede691b 100644 --- a/tests/unit/input_pipeline/olmo_data_grain_resume_test.py +++ b/tests/unit/input_pipeline/olmo_data_grain_resume_test.py @@ -43,6 +43,10 @@ make_olmo_grain_data_loader, ) +import pytest + +pytestmark = pytest.mark.cpu_only + def _write_raw_uint32(tmpdir: str, name: str, values: np.ndarray) -> str: assert values.dtype == np.uint32 and values.ndim == 1 diff --git a/tests/unit/input_pipeline/olmo_data_grain_test.py b/tests/unit/input_pipeline/olmo_data_grain_test.py index 08a493d1c3..168ac57eaa 100644 --- a/tests/unit/input_pipeline/olmo_data_grain_test.py +++ b/tests/unit/input_pipeline/olmo_data_grain_test.py @@ -46,6 +46,10 @@ make_olmo_grain_data_loader, ) +import pytest + +pytestmark = pytest.mark.cpu_only + def _write_raw_uint32(tmpdir: str, name: str, values: np.ndarray) -> str: """Write a 1-D uint32 array as raw binary (no .npy header) — matches AI2.""" From 51362d16a29be874740e228f5bbba652d26409e2 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Sun, 14 Jun 2026 00:54:49 +0000 Subject: [PATCH 42/52] Fix GPU collective ops failure: force early JAX init to prevent context corruption from Qwix/TF --- tests/conftest.py | 25 ++++++------------------- 1 file changed, 6 insertions(+), 19 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index eec4afa225..148a589ee1 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -45,27 +45,14 @@ import os import importlib.util -# Force early JAX initialization on GPU to prevent CUDA context conflicts with TensorFlow/PyTorch. -# If JAX initialization is deferred, TensorFlow/PyTorch (imported during test collection) +# Force early JAX initialization on all platforms to prevent CUDA context conflicts with TensorFlow/PyTorch/Qwix. +# If JAX initialization is deferred, other libraries (imported during test collection) # might initialize CUDA first, causing JAX's subsequent NCCL communicator creation to fail # with 'corrupted comm object detected'. -# Detect GPU environment using standard JAX env vars, GHA runner device types, -# and nvidia-docker visible device markers. -_jax_platforms = os.getenv("JAX_PLATFORMS", "").lower() -_device_type = os.getenv("INPUTS_DEVICE_TYPE", "").lower() -_has_gpu = ( - "cuda" in _jax_platforms - or "gpu" in _jax_platforms - or "cuda" in _device_type - or "gpu" in _device_type - or os.getenv("CUDA_VISIBLE_DEVICES") is not None - or os.getenv("NVIDIA_VISIBLE_DEVICES") is not None -) -if _has_gpu: - try: - _ = jax.devices() - except Exception: # pylint: disable=broad-exception-caught - pass +try: + _ = jax.devices() +except Exception: # pylint: disable=broad-exception-caught + pass # --- Monkeypatch for absl.testing.parameterized --- # Context: Decorating a test method with @parameterized.named_parameters returns a custom From b6529bb7df35f8e11b124cbcaa24aa00245d9268 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Mon, 15 Jun 2026 16:40:58 +0000 Subject: [PATCH 43/52] Update early JAX init logic TAG=agy CONV=3b4d4117-7188-4ee4-bba3-ef97997bd635 --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 148a589ee1..72d72c1632 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,13 +44,13 @@ import jax import os import importlib.util - # Force early JAX initialization on all platforms to prevent CUDA context conflicts with TensorFlow/PyTorch/Qwix. # If JAX initialization is deferred, other libraries (imported during test collection) # might initialize CUDA first, causing JAX's subsequent NCCL communicator creation to fail # with 'corrupted comm object detected'. try: _ = jax.devices() + # Call the internal function if it was previously wrapped, or run raw at the top level except Exception: # pylint: disable=broad-exception-caught pass From a6883c2f7237889294aad42bc760a83bfbbd4202 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Mon, 15 Jun 2026 18:34:36 +0000 Subject: [PATCH 44/52] Fix CPU unit test mock targeting in managed_mldiagnostics_test Fix CPU unit test failures in MaxText by correcting the mock targeting in tests/unit/managed_mldiagnostics_test.py. Patch the module-level variable managed_mldiagnostics.mldiag instead of the class object ManagedMLDiagnostics.mldiag. TAG=agy CONV=c308f588-3466-4c58-a889-2e654d2f2b39 --- tests/unit/managed_mldiagnostics_test.py | 82 ++++++++++++++++++++++++ 1 file changed, 82 insertions(+) create mode 100644 tests/unit/managed_mldiagnostics_test.py diff --git a/tests/unit/managed_mldiagnostics_test.py b/tests/unit/managed_mldiagnostics_test.py new file mode 100644 index 0000000000..2544920028 --- /dev/null +++ b/tests/unit/managed_mldiagnostics_test.py @@ -0,0 +1,82 @@ +# Copyright 2026 Google LLC +# +# 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 +# +# https://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. + +"""Unit tests for ManagedMLDiagnostics.""" + +import unittest +from unittest import mock + +from maxtext.common import managed_mldiagnostics +from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics +import pytest + + +@pytest.mark.cpu_only +class ManagedMLDiagnosticsTest(unittest.TestCase): + # pylint: disable=protected-access + + def setUp(self): + super().setUp() + # Reset singleton instance between tests + ManagedMLDiagnostics._instance = None + + def test_not_enabled_noop(self): + mock_config = mock.MagicMock() + mock_config.managed_mldiagnostics = False + + with mock.patch.object(managed_mldiagnostics.mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(mock_config) + mock_run.assert_not_called() + + def test_enabled_empty_region_passes_none(self): + mock_config = mock.MagicMock() + mock_config.managed_mldiagnostics = True + mock_config.managed_mldiagnostics_region = "" + mock_config.run_name = "test_run" + mock_config.managed_mldiagnostics_run_group = "test_group" + mock_config.managed_mldiagnostics_dir = "gs://test_dir" + mock_config.get_keys.return_value = {"key1": "val1"} + + with mock.patch.object(managed_mldiagnostics.mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(mock_config) + mock_run.assert_called_once_with( + name="test_run", + run_group="test_group", + configs={"key1": "val1"}, + gcs_path="gs://test_dir", + region=None, + ) + + def test_enabled_populated_region_passes_region(self): + mock_config = mock.MagicMock() + mock_config.managed_mldiagnostics = True + mock_config.managed_mldiagnostics_region = "us-east1" + mock_config.run_name = "test_run" + mock_config.managed_mldiagnostics_run_group = "test_group" + mock_config.managed_mldiagnostics_dir = "gs://test_dir" + mock_config.get_keys.return_value = {"key1": "val1"} + + with mock.patch.object(managed_mldiagnostics.mldiag, "machinelearning_run") as mock_run: + ManagedMLDiagnostics(mock_config) + mock_run.assert_called_once_with( + name="test_run", + run_group="test_group", + configs={"key1": "val1"}, + gcs_path="gs://test_dir", + region="us-east1", + ) + + +if __name__ == "__main__": + unittest.main() From 9c2015a523b6ac2d9d5c8f5df2c51170dc7c26b1 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Mon, 15 Jun 2026 21:52:24 +0000 Subject: [PATCH 45/52] Restore NCCL_SOCKET_IFNAME=lo to fix GPU tests --- .github/workflows/run_tests_against_package.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 9bccc54e7a..956cc7b0d9 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -174,9 +174,9 @@ jobs: fi # Configure NCCL for GPU execution - # Removed NCCL_SOCKET_IFNAME=lo as it breaks multi-gpu collective ops in Docker + export NCCL_SOCKET_IFNAME=lo export NCCL_DEBUG=WARN - echo "Set NCCL_DEBUG=WARN for GPU execution." + echo "Set NCCL_DEBUG=WARN and NCCL_SOCKET_IFNAME=lo for GPU execution." fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist From 439d315f927fc297525bffbc5c052320e4da1693 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Wed, 17 Jun 2026 00:39:25 +0000 Subject: [PATCH 46/52] Fix GPU test failures by replacing early jax.devices() with tf.config.set_visible_devices The previous fix attempted to prevent TF/Qwix from grabbing the GPU context by forcing an unconditional jax.devices() call in conftest.py. However, calling jax.devices() at import time locks the XLA client configuration prematurely, which breaks downstream integration tests that dynamically set environment variables like NVTE_FUSED_ATTN or context_parallelism before JAX initialization. This caused those tests to experience mismatched shapes/configurations, leading to a cascade of NCCL invalid argument and corrupted comm errors. This commit replaces the early jax.devices() call with a safer alternative: calling tf.config.set_visible_devices([], 'GPU') directly in conftest.py. This successfully prevents TF (imported transitively by qwix) from allocating the GPU without forcing early JAX initialization, resolving the NCCL failures. TAG=agy CONV=1859e9f1-5466-495f-9e75-e9c7c5b3d4e8 --- .github/workflows/run_tests_against_package.yml | 4 ++-- tests/conftest.py | 10 ++++------ 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 956cc7b0d9..9bccc54e7a 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -174,9 +174,9 @@ jobs: fi # Configure NCCL for GPU execution - export NCCL_SOCKET_IFNAME=lo + # Removed NCCL_SOCKET_IFNAME=lo as it breaks multi-gpu collective ops in Docker export NCCL_DEBUG=WARN - echo "Set NCCL_DEBUG=WARN and NCCL_SOCKET_IFNAME=lo for GPU execution." + echo "Set NCCL_DEBUG=WARN for GPU execution." fi if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist diff --git a/tests/conftest.py b/tests/conftest.py index 72d72c1632..495e9bea5c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -44,13 +44,11 @@ import jax import os import importlib.util -# Force early JAX initialization on all platforms to prevent CUDA context conflicts with TensorFlow/PyTorch/Qwix. -# If JAX initialization is deferred, other libraries (imported during test collection) -# might initialize CUDA first, causing JAX's subsequent NCCL communicator creation to fail -# with 'corrupted comm object detected'. +# Prevent TensorFlow (which might be imported by qwix or other libraries during test collection) +# from grabbing the GPU context and causing JAX NCCL communicator creation to fail. try: - _ = jax.devices() - # Call the internal function if it was previously wrapped, or run raw at the top level + import tensorflow as tf + tf.config.set_visible_devices([], "GPU") except Exception: # pylint: disable=broad-exception-caught pass From 2eceeaa46a5e22e0dde09556116a8d4149c6f0bd Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Wed, 17 Jun 2026 00:52:03 +0000 Subject: [PATCH 47/52] Fix static code-quality checker failures (pylint and formatting) Resolves pyink formatting error in tests/conftest.py introduced by previous commit. Resolves multiple pre-existing pylint errors in types.py, maxengine.py, and attention_op.py that were failing the CodeQuality checks because those files were touched in this branch. TAG=agy CONV=1859e9f1-5466-495f-9e75-e9c7c5b3d4e8 --- src/maxtext/configs/types.py | 2 +- src/maxtext/inference/maxengine/maxengine.py | 2 +- src/maxtext/layers/attention_op.py | 4 ++-- tests/conftest.py | 1 + 4 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index e5bcc9319f..f7b175d1ec 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -2406,7 +2406,7 @@ def _validate_check_vma_is_supported(self): if self.use_ring_of_experts: raise ValueError("check_vma is not yet supported with ring of experts.") _allowed = {"ici_expert_parallelism", "ici_fsdp_parallelism"} - active = [name for name in IciParallelism.model_fields if name not in _allowed and getattr(self, name) != 1] + active = [name for name in IciParallelism.model_fields if name not in _allowed and getattr(self, name) != 1] # pylint: disable=not-an-iterable if active: raise ValueError( f"check_vma=True only supports ici_expert_parallelism and ici_fsdp_parallelism. " diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 189f11044a..95654ec3c5 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -29,7 +29,7 @@ if jax.__version_info__ >= (0, 6, 3): from jax.experimental.layout import Layout as DLL # type: ignore else: - from jax.experimental.layout import DeviceLocalLayout as DLL # type: ignore + from jax.experimental.layout import DeviceLocalLayout as DLL # type: ignore # pylint: disable=no-name-in-module from flax import linen as nn from flax import nnx diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index d1d90a12d8..9a912eb4d1 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -143,7 +143,7 @@ def validate_gpu_flash_attention(sinks: Array | None, record_max_logits: bool) - # TODO(agagik): change splash_attention_mask._ComputableMask to be non protected -class ChunkedCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=protected-access +class ChunkedCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=protected-access,abstract-method """Lazy chunked causal mask. Attention is causal within each chunk (0, K), (K, 2K), (2K, 3K), ... tokens @@ -2157,7 +2157,7 @@ def __call__( # pylint: disable=protected-access -class LoadBalancedCausalMask(splash_attention_mask._ComputableMask): +class LoadBalancedCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=abstract-method """Lazy causal mask, prevents the model from attending to future tokens. Attributes: diff --git a/tests/conftest.py b/tests/conftest.py index 495e9bea5c..b0eca11604 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -48,6 +48,7 @@ # from grabbing the GPU context and causing JAX NCCL communicator creation to fail. try: import tensorflow as tf + tf.config.set_visible_devices([], "GPU") except Exception: # pylint: disable=broad-exception-caught pass From 341018b9d2f3a234006e714442684a380bbfe1df Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 25 Jun 2026 05:30:55 +0000 Subject: [PATCH 48/52] Fix nnx wrapper method binding, add vocab tiling intermediates, and support recursive intermediate retrieval --- src/maxtext/layers/nnx_wrappers.py | 8 +- src/maxtext/models/models.py | 2 + src/maxtext/utils/maxtext_utils.py | 82 +++++++++++++++--- tests/unit/maxtext_utils_test.py | 129 +++++++++++++++++++++++------ tests/unit/tiling_test.py | 70 ++++++++++++++++ 5 files changed, 252 insertions(+), 39 deletions(-) diff --git a/src/maxtext/layers/nnx_wrappers.py b/src/maxtext/layers/nnx_wrappers.py index b483649c9e..8e04e40939 100644 --- a/src/maxtext/layers/nnx_wrappers.py +++ b/src/maxtext/layers/nnx_wrappers.py @@ -14,6 +14,7 @@ """NNX <> Linen interoperability.""" +import types from functools import partial import typing as tp from typing import Any @@ -526,9 +527,9 @@ def __getattr__(self, name: str): return self.kwargs[name] maybe_method = getattr(self.nnx_class, name, None) if callable(maybe_method): - method = partial(self.__call__, nnx_method=maybe_method) - method.__self__ = self - return method + def unbound_func(instance, *args, **kwargs): + return instance.__call__(*args, nnx_method=maybe_method, **kwargs) + return types.MethodType(unbound_func, self) return super().__getattribute__(name) def _update_variables(self, module): @@ -545,6 +546,7 @@ def _update_variables(self, module): collection_flat_state[collection] = [] collection_flat_state[collection].append((path, leaf)) + # update linen variables for collection, flat_state in collection_flat_state.items(): if self.is_mutable_collection(collection): diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index 20ee4d1263..4785fb1ee6 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -533,6 +533,8 @@ def __call__( mutable_collections.append("intermediates") if self.config.load_balance_loss_weight > 0.0 and "intermediates" not in mutable_collections: mutable_collections.append("intermediates") + if self.config.num_vocab_tiling > 1 and "intermediates" not in mutable_collections: + mutable_collections.append("intermediates") if self.config.pure_nnx_decoder: logits, hidden_state, kv_caches = self.decoder( diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 16b022c3a4..efb2bf795c 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -1281,10 +1281,39 @@ def collect_intermediates_by_suffix(intermediate_outputs, *suffix_keys: str) -> return values -def get_intermediate_value(model, nested_key, default=None, clear=False): +def _find_and_remove_intermediates(state, suffix, clear=False): + """Recursively finds intermediate values matching suffix. + + If clear=True, removes them from the state. + Returns a list of (path, variable) tuples. """ - Retrieves an intermediate value from an NNX model. This functions has context about - where the intermediate value is located. + results = [] + + def _traverse(current_state, current_path): + keys_to_delete = [] + for k, v in list(current_state.items()): + new_path = current_path + (k,) + if isinstance(v, nnx.Intermediate): + if len(new_path) >= len(suffix) and new_path[-len(suffix):] == suffix: + results.append((new_path, v)) + if clear: + keys_to_delete.append(k) + elif isinstance(v, (nnx.State, dict)): + _traverse(v, new_path) + if clear and not v: + keys_to_delete.append(k) + + for k in keys_to_delete: + del current_state[k] + + _traverse(state, ()) + return results + + +def get_intermediate_value(model, nested_key, default=None, clear=False): + """Retrieves an intermediate value from an NNX model. + + This functions has context about where the intermediate value is located. Args: model: The NNX model. @@ -1295,18 +1324,51 @@ def get_intermediate_value(model, nested_key, default=None, clear=False): Returns: The value associated with the nested key, or the default value if not found. """ - intermediate_value = default match nested_key: case "out_projection_activations": - if nested_key in model.decoder.layers["self_attention"]: - intermediate_value = model.decoder.layers["self_attention"][nested_key].get_value()[-1] - if clear: - del model.decoder.layers["self_attention"][nested_key] + suffixes = [ + ("self_attention", "out_projection_activations"), + ("GptOssAttention", "out_projection_activations"), + ] case _: - # Default case to handle any unknown nested keys raise ValueError(f"Incorrect nested_key: {nested_key}") - return intermediate_value + # Pop all intermediates to safely inspect and potentially clear them + intermediates = nnx.pop(model, nnx.Intermediate) + + found = [] + for suffix in suffixes: + found = _find_and_remove_intermediates(intermediates, suffix, clear=clear) + if found: + break + + # Put back the remaining intermediates + nnx.update(model, intermediates) + + if not found: + return default + + # Helper key function to sort paths numerically if indices are present + def path_sort_key(item): + path = item[0] + def _to_int_if_possible(val): + if isinstance(val, int): + return val + if isinstance(val, str) and val.isdigit(): + return int(val) + return val + return tuple(_to_int_if_possible(x) for x in path) + + found.sort(key=path_sort_key) + + values = [var.get_value()[-1] for path, var in found] + + if len(values) > 1: + # Multiple layers (sequential), stack them + return jnp.stack(values, axis=0) + else: + # Single layer (scanned or just 1 layer), return directly + return values[0] def update_state_param(state, target_path, value): diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 3d4e983281..61d7286cae 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -95,45 +95,122 @@ def test_fp8_stats_not_clipped_but_others_are(self): ) ) +class MockAttention(nnx.Module): + def __init__(self, rngs): + pass + def __call__(self, x, sow_val=None): + if sow_val is not None: + self.sow(nnx.Intermediate, "out_projection_activations", sow_val) + return x -class TestIntermediateValueRetrieval(unittest.TestCase): - """test class for IntermediateValueRetrieval""" - - def setUp(self): - self.mock_model = MagicMock(name="Transformer") +class MockDecoderLayer(nnx.Module): + def __init__(self, rngs, attention_name="self_attention"): + if attention_name == "self_attention": + self.self_attention = MockAttention(rngs) + elif attention_name == "GptOssAttention": + self.GptOssAttention = MockAttention(rngs) + self.attention_name = attention_name + + def __call__(self, x, sow_val=None): + attention = getattr(self, self.attention_name) + return attention(x, sow_val) + +class MockDecoderSequential(nnx.Module): + def __init__(self, rngs, attention_name="self_attention"): + self.layers = nnx.List([ + MockDecoderLayer(rngs, attention_name), + MockDecoderLayer(rngs, attention_name), + ]) + + def __call__(self, x, sow_vals=None): + if sow_vals is None: + sow_vals = [None, None] + out = x + for layer, val in zip(self.layers, sow_vals): + out = layer(out, val) + return out + +class MockDecoderScanned(nnx.Module): + def __init__(self, rngs, attention_name="self_attention"): + self.layers = MockDecoderLayer(rngs, attention_name) + self.attention_name = attention_name + + def __call__(self, x, sow_val=None): + if sow_val is not None: + attention = getattr(self.layers, self.attention_name) + attention.sow(nnx.Intermediate, "out_projection_activations", sow_val) + return x - # 2. Create the Decoder Mock - self.mock_decoder = MagicMock(name="Decoder") - self.mock_model.decoder = self.mock_decoder - self.mock_layers = {} - self.mock_model.decoder.layers = self.mock_layers - self.self_attention = {} - self.mock_layers["self_attention"] = self.self_attention +class MockTransformer(nnx.Module): + def __init__(self, decoder_type, rngs, attention_name="self_attention"): + if decoder_type == "sequential": + self.decoder = MockDecoderSequential(rngs, attention_name) + elif decoder_type == "scanned": + self.decoder = MockDecoderScanned(rngs, attention_name) - def test_valid_intermediate_key(self): - expected_sowed_data = [0.1, 0.5, 0.9] - mock_sowed_variable = Mock(name="out_projection_activations") - mock_sowed_variable.get_value.return_value = (expected_sowed_data,) + def __call__(self, x, sow_vals=None): + return self.decoder(x, sow_vals) - self.mock_decoder.layers["self_attention"]["out_projection_activations"] = mock_sowed_variable - result = maxtext_utils.get_intermediate_value(self.mock_model, "out_projection_activations") +class TestIntermediateValueRetrieval(unittest.TestCase): + """test class for IntermediateValueRetrieval""" - self.assertEqual(result, expected_sowed_data) + def test_valid_intermediate_key_sequential(self): + rngs = nnx.Rngs(0) + model = MockTransformer("sequential", rngs) + x = jnp.ones((1, 2)) + sow_vals = [jnp.array([0.1]), jnp.array([0.5])] + model(x, sow_vals) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations") + expected = jnp.stack(sow_vals, axis=0) + self.assertTrue(jnp.allclose(result, expected)) + + def test_valid_intermediate_key_scanned(self): + rngs = nnx.Rngs(0) + model = MockTransformer("scanned", rngs) + x = jnp.ones((1, 2)) + sow_val = jnp.array([[0.1], [0.5]]) + model(x, sow_val) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations") + self.assertTrue(jnp.allclose(result, sow_val)) + + def test_valid_intermediate_key_sequential_gpt_oss(self): + rngs = nnx.Rngs(0) + model = MockTransformer("sequential", rngs, attention_name="GptOssAttention") + x = jnp.ones((1, 2)) + sow_vals = [jnp.array([0.1]), jnp.array([0.5])] + model(x, sow_vals) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations") + expected = jnp.stack(sow_vals, axis=0) + self.assertTrue(jnp.allclose(result, expected)) + + def test_valid_intermediate_key_scanned_gpt_oss(self): + rngs = nnx.Rngs(0) + model = MockTransformer("scanned", rngs, attention_name="GptOssAttention") + x = jnp.ones((1, 2)) + sow_val = jnp.array([[0.1], [0.5]]) + model(x, sow_val) + + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations") + self.assertTrue(jnp.allclose(result, sow_val)) def test_returns_default_if_sow_did_not_happen(self): - """ - Simulate a scenario where the model ran, but this specific key - was NOT sowed (or the layer was skipped). - """ - - result = maxtext_utils.get_intermediate_value(self.mock_model, "out_projection_activations", default="MyDefault") + rngs = nnx.Rngs(0) + model = MockTransformer("sequential", rngs) + x = jnp.ones((1, 2)) + model(x, None) + result = maxtext_utils.get_intermediate_value(model, "out_projection_activations", default="MyDefault") self.assertEqual(result, "MyDefault") def test_unknown_key_raises_value_error(self): + rngs = nnx.Rngs(0) + model = MockTransformer("sequential", rngs) with self.assertRaises(ValueError) as cm: - maxtext_utils.get_intermediate_value(self.mock_model, "some_random_layer_name") + maxtext_utils.get_intermediate_value(model, "some_random_layer_name") self.assertEqual(str(cm.exception), "Incorrect nested_key: some_random_layer_name") diff --git a/tests/unit/tiling_test.py b/tests/unit/tiling_test.py index 899c1227e4..0473e797e8 100644 --- a/tests/unit/tiling_test.py +++ b/tests/unit/tiling_test.py @@ -268,6 +268,76 @@ def test_vocab_tiling_gradient_with_z_loss(self): "Gradients do not match for vocab tiling when z-loss is enabled.", ) + def test_vocab_tiling_gradient_with_z_loss_nnx(self): + """ + Tests loss and gradient correctness when z-loss is enabled, comparing + standard computation vs. vocabulary tiling computation with NNX enabled. + """ + cfg_non_tiling = pyconfig.initialize( + self.base_config, + run_name="grad_test_z_loss_no_tiling_nnx", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=1, + z_loss_multiplier=1e-4, # Enable z-loss + enable_nnx=True, + ) + quant_non_tiling = quantizations.configure_quantization(cfg_non_tiling) + devices_array_non_tiling = maxtext_utils.create_device_mesh(cfg_non_tiling) + mesh_non_tiling = Mesh(devices_array_non_tiling, cfg_non_tiling.mesh_axes) + model_non_tiling = models.transformer_as_linen( + cfg_non_tiling, mesh=mesh_non_tiling, quant=quant_non_tiling, model_mode=MODEL_MODE_TRAIN + ) + + rng_model, rng_targets = jax.random.split(self.rng) + + params = model_non_tiling.init( + {"params": rng_model, "dropout": rng_model}, + self.dummy_inputs, + self.dummy_inputs, + ) + + data = { + "targets": jax.random.randint(rng_targets, (self.batch_size, self.seq_len), 0, cfg_non_tiling.vocab_size), + "targets_segmentation": jnp.ones((self.batch_size, self.seq_len)), + } + + loss_non_tiling, grads_non_tiling = self.get_grads(cfg_non_tiling, params, data) + + cfg_tiling = pyconfig.initialize( + self.base_config, + run_name="grad_test_z_loss_with_tiling_nnx", + enable_checkpointing=False, + enable_dropout=False, + max_target_length=self.seq_len, + per_device_batch_size=self.batch_size, + logits_via_embedding=False, + base_num_decoder_layers=0, + dtype="float32", + matmul_precision="high", + num_vocab_tiling=4, + z_loss_multiplier=1e-4, # Enable z-loss + enable_nnx=True, + ) + loss_tiling, grads_tiling = self.get_grads(cfg_tiling, params, data) + + # Loss correctness test + assert jnp.allclose(loss_non_tiling, loss_tiling, rtol=self.rtol), "Losses do not match when z-loss is enabled (NNX)." + + # Gradient correctness test + self.assert_pytrees_all_close( + grads_non_tiling, + grads_tiling, + "Gradients do not match for vocab tiling when z-loss is enabled (NNX).", + ) + + @pytest.mark.tpu_only def test_vocab_tiling_nnx_loss(self): """ From b1685a4cf4b4e170c743c8ee141dcdffd6f795b5 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 25 Jun 2026 05:41:09 +0000 Subject: [PATCH 49/52] Fix AqtQuantization reference in deepseek4.py --- src/maxtext/models/deepseek4.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/maxtext/models/deepseek4.py b/src/maxtext/models/deepseek4.py index c259db8c66..25eca866f9 100644 --- a/src/maxtext/models/deepseek4.py +++ b/src/maxtext/models/deepseek4.py @@ -56,7 +56,7 @@ def __init__( model_mode: str, mesh: Mesh, rngs: nnx.Rngs, - quant: Optional[quantizations.AqtQuantization] = None, + quant: Optional[quantizations.Quantization] = None, layer_idx: int = -1, compress_ratio: Optional[int] = None, is_hash_routing: Optional[bool] = None, @@ -184,7 +184,7 @@ def __init__( mesh: Mesh, model_mode: str, rngs: nnx.Rngs, - quant: None | quantizations.AqtQuantization = None, + quant: None | quantizations.Quantization = None, ): self.config = config self.mesh = mesh From 972e6f5f40ea1c0b7ede095966f7792b7acea4cf Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 25 Jun 2026 05:47:03 +0000 Subject: [PATCH 50/52] Fix pylint error: add missing traversals import --- tests/unit/quantizations_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/quantizations_test.py b/tests/unit/quantizations_test.py index 68f2387cf6..5aa252c9aa 100644 --- a/tests/unit/quantizations_test.py +++ b/tests/unit/quantizations_test.py @@ -22,6 +22,7 @@ from maxtext.configs import pyconfig from maxtext.common.common_types import DECODING_ACTIVE_SEQUENCE_INDICATOR from flax import nnx +from flax.nnx import traversals from maxtext.layers import moe from maxtext.layers import linears from maxtext.layers import quantizations From 8df6d3101ecc6bcd1f392a04fa21f0f863537e1b Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 25 Jun 2026 14:54:15 +0000 Subject: [PATCH 51/52] Fix pylint missing docstrings and reformat files --- .../aqt_deprecation_phase_2_execution_plan.md | 187 ++++++++++++++ docs/uxr_usability_test_plan.md | 235 ++++++++++++++++++ resolve_conflict.py | 23 ++ src/maxtext/inference/maxengine/maxengine.py | 8 +- src/maxtext/layers/nnx_wrappers.py | 3 +- src/maxtext/utils/maxtext_utils.py | 12 +- tests/unit/maxtext_utils_test.py | 32 ++- tests/unit/tiling_test.py | 1 - 8 files changed, 478 insertions(+), 23 deletions(-) create mode 100644 docs/aqt_deprecation_phase_2_execution_plan.md create mode 100644 docs/uxr_usability_test_plan.md create mode 100644 resolve_conflict.py diff --git a/docs/aqt_deprecation_phase_2_execution_plan.md b/docs/aqt_deprecation_phase_2_execution_plan.md new file mode 100644 index 0000000000..86be007916 --- /dev/null +++ b/docs/aqt_deprecation_phase_2_execution_plan.md @@ -0,0 +1,187 @@ +# MaxText AQT Deprecation Phase #2: Technical Execution Plan + +This document provides a comprehensive, highly technical, and definitive execution plan for Phase #2 of the Accurate Quantization Training (AQT) deprecation within the MaxText repository. + +Our primary goal is to **completely strip out the legacy AQT-specific code, configurations, and dependencies** while **maintaining absolute architectural generality** for modern quantization backends (such as Qwix, native FP8, and TransformerEngine). + +--- + +## 1. Architectural Code Review & Impact Assessment + +We have performed a deep-dive architectural review of the AQT footprint across the MaxText repository. Below is the precise mapping of AQT dependencies, configuration flags, and the "blast radius" of their removal. + +```mermaid +graph TD + subgraph Configuration + types_py["configs/types.py"] --> |Defines| QuantType["QuantizationType Enum"] + types_py --> |Defines| QuantConfig["Quantization Class"] + end + + subgraph Core Quantization + quant_py["layers/quantizations.py"] --> |Imports| AQT_v2["aqt.jax.v2"] + quant_py --> |Implements| AqtQuant["AqtQuantization"] + quant_py --> |Implements| MemoryUtils["remove_quantized_params"] + end + + subgraph Layers & Models + linears_py["layers/linears.py"] --> |Calls| AqtQuant + models_py["models/models.py"] --> |Type Hints| AqtQuant + decoders_py["layers/decoders.py"] --> |Type Hints| AqtQuant + end + + subgraph Tooling & Inference + lw_quant["utils/layerwise_quantization.py"] --> |Imports| AQT_v2 + maxengine["inference/maxengine/maxengine.py"] --> |Calls| MemoryUtils + end + + AqtQuant -.-> |Injects| linears_py + AqtQuant -.-> |Injects| models_py + AqtQuant -.-> |Injects| decoders_py +``` + +### A. Specific Modules, Files, and Configuration Flags to Modify/Remove + +1. **External Library Dependencies:** + * **File:** [requirements.txt](file:///usr/local/google/home/sarunsingla/maxtext/src/dependencies/requirements/requirements.txt#L2), [base_requirements/requirements.txt](file:///usr/local/google/home/sarunsingla/maxtext/src/dependencies/requirements/base_requirements/requirements.txt#L2), and generated cuda/tpu lockfiles. + * **Action:** Remove the `aqtp` package dependency (which installs `aqt` in Python). +2. **Configuration Flags & Enums:** + * **File:** [types.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py) + * **Flags to Remove:** + * `replicate_quant_scale`: AQT-specific flag used to replicate scales across mesh axes to avoid inefficient XLA fusions. + * `quant_cfg_path`: Path for `intmp` (AQT mixed precision) configurations. + * **Quantization Types to Prune:** Remove `INTMP = "intmp"`, `aqt_fp8`, and `aqt_fp8_full` from [QuantizationType](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py#L83) enum and config parsing. + * **Flag to Deprecate/Modify:** + * `use_qwix_quantization`: Currently defaults to `False`. For Phase 2, we will set this default to `True` or completely remove it, making Qwix/native FP8 the standard quantization pipeline. +3. **Core Quantization Library:** + * **File:** [quantizations.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py) + * **Action:** + * Remove all `aqt.jax.v2` and `aqt_flax` imports. + * Delete the [AqtQuantization](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py#L119) dataclass. + * Delete AQT helper functions: `_tiling_fn`, `_rhs_axis_metadata_wrapper`, `_build_const_scale_config`, `_build_per_tensor_config`, `_get_int8_quant_config`, `_get_aqt_fp8_quant_config`, `_get_aqt_fp8_default_config`, `_dot_general_make`, `_get_default_mp_config`, `_get_mixed_precision_quant_config`. + * Delete AQT memory optimization utilities: `match_aqt_and_unquantized_param`, `_get_aqt_key_paths`, and [remove_quantized_params](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py#L695). +4. **Inference Engine Integration:** + * **File:** [maxengine.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/inference/maxengine/maxengine.py#L524) + * **Action:** Remove the call to `quantizations.remove_quantized_params` and the handling of the `"aqt"` variable collection. +5. **Layer-wise Quantization Tooling:** + * **File:** [layerwise_quantization.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/utils/layerwise_quantization.py) + * **Action:** This utility is deeply coupled to AQT's `QTensor` and `AqtDotGeneral` structures. We will completely deprecate this file or refactor it to target Qwix/native FP8. + +### B. Blast Radiuses & Critical Decoupling +* **The Quantization Param Type Hint:** Over 20 files (including [models.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/models/models.py#L36), [linears.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/linears.py#L37), [decoders.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/decoders.py#L40), and [moe.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/moe.py#L192)) import and type-hint the `quant` parameter using `AqtQuantization as Quant`. + * *Mitigation:* We must **not** delete the base class [Quantization](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py#L55). Instead, we will import and type-hint with `Quantization as Quant`. This keeps layer signatures clean and compatible with Qwix, native FP8, and TransformerEngine. +* **Layer Computations:** In [linears.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/linears.py#L70), `_compute_dot_general` and `_compute_dot_general_nnx` invoke `quant.dot_general_cls()`. + * *Mitigation:* As long as other quantization backends inherit from the base `Quantization` class and implement `dot_general_cls()`, their execution remains completely unaffected. + +--- + +## 2. Phase #2 Deprecation Plan & Scope of Changes + +We will execute the AQT removal using a structured, phased engineering sprint. + +### Step 1: Remove Third-Party Dependency & Clean Configs +1. Prune `aqtp` from all requirements lockfiles. +2. Edit [types.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py) to remove `replicate_quant_scale` and `quant_cfg_path` from the [Quantization](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py#L423) configuration class. +3. Remove the AQT warning check from `set_derived_and_validate_values` in [types.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py#L2574-L2580). +4. Throw a compile-time `ValueError` if a user attempts to run quantization with `use_qwix_quantization=False` and the quantization type is not native FP8 or TransformerEngine. + +### Step 2: Refactor `layers/quantizations.py` +Modify [quantizations.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py) to remove all AQT references: + +```diff +-from aqt.jax.v2 import config as aqt_config +-from aqt.jax.v2 import aqt_tensor +-from aqt.jax.v2.flax import aqt_flax +-from aqt.jax.v2 import tiled_dot_general +-from aqt.jax.v2 import calibration + +... + +-@dataclass +-class AqtQuantization: +- """Configures AQT quantization github.com/google/aqt.""" +- ... +- def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): +- ... +- def einsum(self, mesh_axes: Tuple[str, ...] = ()): +- ... + +... + +-def remove_quantized_params(params, aqt_vars): +- ... +``` + +### Step 3: Generalize Layer and Model Type-Hints +In every model and layer file, update the imports and type hints: + +```diff +-from maxtext.layers.quantizations import AqtQuantization as Quant ++from maxtext.layers.quantizations import Quantization as Quant +``` + +This change preserves the signature of all layers: +```python +class DenseGeneral(nnx.Module): + def __init__( + self, + ... + quant: None | Quant = None, + ... + ): +``` + +### Step 4: Prune Inference & Tooling Hooks +1. In [maxengine.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/inference/maxengine/maxengine.py#L524), remove AQT parameter extraction and the `remove_quantized_params` call. +2. Delete or refactor [layerwise_quantization.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/utils/layerwise_quantization.py) to utilize Qwix quantization rules instead of AQT's `QTensor` structures. + +--- + +## 3. End-to-End (E2E) Testing Strategy + +To guarantee that this cleanup introduces zero regressions in model correctness, hardware utilization, or training stability, we will execute the following multi-tier verification matrix. + +``` + [Unit Tests] [Integration Tests] [E2E Training Runs] + (quantizations_test.py) (maxengine_test.py) (Llama2-7B / Gemma2-9B) + │ │ │ + ├─► Prune AQT tests ├─► Test Qwix/FP8 ├─► Verify Convergence + └─► Verify Qwix/FP8 └─► Verify serving latency └─► Check TPU/GPU MFU +``` + +### A. Unit & Integration Tests +* **File to Refactor:** [quantizations_test.py](file:///usr/local/google/home/sarunsingla/maxtext/tests/unit/quantizations_test.py) + * **Action:** + * Remove `QuantTestModule`'s explicit dependency on `aqt_flax.AqtDotGeneral` and `aqt_flax.AqtEinsum`. + * Prune AQT-specific test cases: `test_aqt_quantization`, `test_mixed_precision_*`, and `test_remove_quantized_params`. + * Enhance existing Qwix and native FP8 tests (`test_int8_quantization`, `test_fp8_quantization`, `test_fp8_full_quantization`) to ensure full coverage of the remaining quantization pathways. +* **Other Unit Tests:** Run `tiling_test.py`, `gpt3_test.py`, `maxtext_utils_test.py`, and `model_test.py` to ensure that standard unquantized runs are completely unaffected. + +### B. Regression Testing (Performance & Correctness) +* **Convergence Verification:** Run a Llama2-7B training run for 1000 steps using BF16 (unquantized) and Qwix FP8. Compare the loss curves before and after the PR to verify mathematical equivalence. +* **Throughput & Utilization:** Monitor Step Time (ms) and Model Flops Utilization (MFU) on TPU v5e/v6e and NVIDIA H100 GPUs. The elimination of AQT code paths must result in **zero throughput regression** and **zero HBM memory overhead increases** for standard models. + +### C. E2E Training Scale Verification +We will run E2E training workloads to validate the entire JIT compilation and execution pipeline: +1. **Workload A (TPU Scale):** Llama2-7B training on a 16-device TPU v6e slice using Qwix FP8 (`quantization=fp8_full` and `use_qwix_quantization=True`). +2. **Workload B (GPU Scale):** Gemma2-9B training on an 8-GPU H100 node using native FP8 (`quantization=fp8` and `use_qwix_quantization=True`). +3. **Workload C (MoE Scale):** DeepSeek-V3 MoE model initialized and executed using TransformerEngine quantization (`quantization=te_fp8_delayedscaling`). + +--- + +## 4. Anticipated Edge Cases, Bug Fixing, and Risk Mitigation + +During this infrastructure migration, we anticipate 3 major technical hurdles. Below are the exact engineering remedies for each. + +### Hurdle 1: Legacy Checkpoint Loading Failure (Structure Mismatch) +* **The Issue:** Legacy Orbax checkpoints saved using AQT quantization contain AQT-specific variables (`qrhs` and `frozen` QTensor structures). Attempting to restore these checkpoints into the new AQT-free model structure will cause Orbax to throw a structural mismatch error because the model no longer defines those variables. +* **The Remedy:** + * Implement an explicit warning or error in [model_creation_utils.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/utils/model_creation_utils.py#L334) during the checkpoint restoration phase. + * Provide a standalone migration script (`maxtext/checkpoint_conversion/remove_aqt_vars.py`) that loads a legacy AQT checkpoint, strips out the AQT-specific variables (converting them back to unquantized parameters or Qwix-compatible weights), and saves a clean Orbax checkpoint. + +### Hurdle 2: Shape Mismatches in Grouped MatMul (GMM) / MoE Kernels +* **The Issue:** AQT deprecation could inadvertently break how tiling dimensions and contraction axes are set up in MoE layers, leading to shape mismatches in custom Pallas kernels (like GMM in [ops.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/kernels/megablox/ops.py)). +* **The Remedy:** Ensure that the generic `Quantization` interface correctly plumbs the tile sizes and contraction axes to `DenseGeneral` and `MlpBlock` in [linears.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/linears.py#L188). Add strict shape validation assertions in [moe.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/moe.py) to verify that the intermediate sharding matches the GMM kernel expectations before launching the kernel. + +### Hurdle 3: Compile-time Tracing Failures in NNX Bridge +* **The Issue:** The NNX-to-Linen bridge ([nnx_wrappers.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/nnx_wrappers.py)) traces the module during `lazy_init` and expects certain mutable collections (like `"aqt"`). If these collections are missing or empty, JAX might throw a tracing error. +* **The Remedy:** In [linears.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/linears.py#L97), update `_compute_dot_general_nnx` to only include `"aqt"` in the `mutable` list if the quantization backend actually registers that collection. For Qwix and native FP8, ensure they use their respective state collections (like `"params"` or `"intermediates"`), preventing JAX tracing crashes. diff --git a/docs/uxr_usability_test_plan.md b/docs/uxr_usability_test_plan.md new file mode 100644 index 0000000000..91e03550cd --- /dev/null +++ b/docs/uxr_usability_test_plan.md @@ -0,0 +1,235 @@ +# UXR Usability Analysis & Ecosystem-Aligned Corner Test Plan for MaxText + +This document presents an analysis of user-facing pain points, usability gaps, and cross-product dependency friction modes in MaxText. It proposes a concrete, phased implementation plan for introducing **specialized corner test suites** and an **integrated cross-product CI pipeline** designed to prevent regressions, improve developer velocity, and ensure stable training and inference runs on large-scale TPU/GPU topologies. + +--- + +## 1. Integrated UXR & Ecosystem Friction Analysis + +MaxText operates within a multi-layered ecosystem of high-performance libraries for training, inference, orchestration, and job submission. Because dependencies release on independent cycles, uncoordinated updates frequently introduce regressions. + +The following table details usability gaps, environment friction points (MT-01 to MT-19), and the proposed testing/CI remedies to shift validation left: + +| ID | Friction / Usability Pain Point | Technical Root Cause & Dependencies | Diagnostic / Impact | Proposed Testing Remedy | +| :--- | :--- | :--- | :--- | :--- | +| **MT-01** | VM restarts wipe newly installed XPK elements. | State non-persistence across VM restarts. | Users lose environment setup, forcing rebuilds. | **Suite F: VM State Persistence Test** (Asserts configuration and tool persistence across simulated restarts). | +| **MT-02** | Stacking post-training packages fails to override older library versions. | Installation conflicts with pre-existing packages (specifically `tunix`). | Opaque Python keyword errors that halt post-training setups. | **Suite E: Environment Package Isolation Audit** (Validates clean overrides and environment resets). | +| **MT-03** | GRPO/SFT workflows require building full Docker images from source. | Absence of pre-built, modular release images and clear build progress. | 10+ minute delay per execution attempt with no progress feedback or build time estimates. | **Suite F: Modular Docker & Pre-built Release Validation** (Verifies compatibility of pre-built images and build timeline output). | +| **MT-04** | Unused heavy dependencies (vLLM, tpu-common) are installed during SFT. | Non-modular Dockerfile configurations. | Unnecessary image compilation overhead (adds 2+ minutes per build). | **Suite F: Lightweight Build Pathway Test** (Validates dependency footprint optimization in SFT Dockerfiles). | +| **MT-05** | Docker upload runner script fails silently. | Opaque utility runner execution without logs or stdout capture. | Users unable to self-diagnose; requires manual code edits to trace. | **Suite F: Verbose Log & Exit Code Auditor** (Enforces standard error logging and explicit exit codes). | +| **MT-06** | Documentation interchanges critical environment variables. | Inconsistent naming convention in tutorials (e.g., `MODEL` vs `MODEL_NAME`). | Continuous user copy-paste failures and runtime syntax errors. | **Suite G: Documentation Linting Suite** (Automated markdown command syntax and variable validation). | +| **MT-07** | Tutorial configurations specify non-instruct models with instruct templates. | Mismatched model weight configs and runtime chat templates. | Immediate execution crashes when running tutorials. | **Suite G: Model Configuration Auditor** (Validates matching characteristics between configs and code). | +| **MT-08** | Prerequisite installation steps are buried or out-of-order in explanatory prose. | Suboptimal layout structure and sequencing in user tutorials. | Users skim text, miss vital setup commands (e.g., installing `uv`), or try to run utilities before they are installed. | **Suite G: Document Structure Auditor** (Enforces separation and logical order of execution steps). | +| **MT-09** | XPK commands require redundant `--zone` flags. | CLI does not inherit global `gcloud` configurations. | Redundant input overhead and minor command-line syntax friction. | **Suite G: CLI Argument Inheritance Test** (Ensures CLI tools respect global fallback configurations). | +| **MT-10** | TPU cluster nomenclature is inconsistent across platforms. | Discrepant naming schemas (e.g., Console `v6e-8` vs GKE `ct6e-standard-4t`). | Hardware configuration confusion during cluster mapping. | **Suite G: Naming Mapping Validator** (Checks correctness of taxonomy mapping utilities). | +| **MT-11** | Copied multi-line code blocks inject hidden newline characters. | Format copying issues in public documentation hub. | Compound execution and string parsing errors in terminal sessions. | **Suite G: Copy-Paste Character Sanitizer** (Tests copy-paste text scrubbing on documentation assets). | +| **MT-12** | Ambiguity regarding git cloning and install choices. | Documentation fails to clarify when steps are mutually exclusive or mandatory. | Users perform redundant setup steps (e.g., installing both vLLM and TPU branches) or miss foundational repository checkouts. | **Suite G: Prerequisite Path Validator** (Verifies dependency branches are distinct and explicitly mutually exclusive). | +| **MT-13** | Missing critical setup templates and commands. | Documentation omits fundamental infrastructure preparation steps (e.g., GCS bucket creation for SFT output). | Users are blocked or forced to construct cloud setup commands manually from scratch. | **Suite G: Command Template Auditor** (Verifies documentation contains copy-pastable templates for all prerequisites). | +| **MT-14** | Path assumptions cause execution errors. | Scripts assume a specific working directory without validating the user's path context. | Users run commands from non-root paths and receive path errors (e.g., with `docker_upload_runner.sh`). | **Suite G: Directory-Agnostic Runner Test** (Validates that utility runners execute successfully regardless of invocation CWD). | +| **MT-15** | Context switching across fragmented documentation silos. | Information is scattered across MaxText GitHub, Cloud, and external libraries. | High cognitive load and slower setup times as users hop between distinct doc ecosystems. | **Suite G: Cross-Doc Link Integrity Suite** (Validates correctness and co-location of essential setup links). | +| **MT-16** | Divergent tutorials for expanding post-training techniques. | Lack of centralized architectural strategy for post-training guides (GRPO, Distillation, PPO, DPO). | Maintenance burden increases and documentation diverges as more techniques are added. | **Suite G: Post-Training Documentation Matrix** (Ensures unified branching structure for post-training guides). | +| **MT-17** | XPK commands lack clear env variable configuration guidance. | Setup scripts fail to document or auto-populate necessary cluster and workload variables. | Users must manually search Cloud Console to find details like TPU type or workload ID to fill XPK parameters. | **Suite F: XPK Environment Variable Audit** (Validates existence of helper env scripts and config parsers). | +| **MT-18** | Incompatible TPU profiles in notebooks. | Colab notebooks contain default links with unsupported or unavailable TPU types (e.g., requesting v5p without access). | Immediate failure during first cell execution in Colab, blocking users. | **Suite E: Notebook Hardware Compatibility Test** (Validates TPU availability matches notebook requirements). | +| **MT-19** | Package manager and SSH command failures on locked VMs. | Standard tools (`sudo apt`) are missing or standard SSH structures fail in stripped environments. | Users cannot update libraries or connect to VMs, leading to absolute blockages. | **Suite F: Base Image Tooling Sanity Test** (Checks that minimal required CLI commands are present and functional). | +| **SEC-01**| Project Editor roles lack write access to the Artifact Registry. | Missing IAM permission guidance during Docker image uploads. | Users blocked immediately on image push, forcing manual IAM fixes. | **Suite F: Pre-Execution IAM Validator** (Checks Registry Writer access before starting upload scripts). | +| **SEC-02**| Missing Docker authentication instructions. | Guide omits required configuration commands (e.g., `gcloud auth configure-docker`). | Silent permission denied errors at the final image upload step. | **Suite F: Pre-Execution Auth Validator** (Validates docker auth config prior to running build actions). | +| **SEC-03**| Opaque TPU VM service account mapping. | Documentation lacks clear guidelines on VM-attached service accounts vs personal credentials. | Standard workflows fail on permission checks, requiring manual troubleshooting. | **Suite F: Service Account IAM Auditor** (Verifies attached TPU VM service accounts possess the required GCS/Registry scopes). | +| **UXR-01**| Late-Stage Configuration Failures. | No ahead-of-time (AOT) configuration validation. | Wasted TPU resource hours; errors caught hours into compilation or step runs. | **Suite A: Static Config & Topology Validator** (AOT checks for mesh dimension alignment and memory). | +| **UXR-02**| Restoration Hangs & Resharding OOMs. | Mesh topology mismatches (e.g., 2D to 3D) during checkpoint restoration. | Silent compilation hangs or abrupt OOM errors during restoration. | **Suite B: Checkpoint Resharding Suite** (Mock-reshard testing across virtual topologies). | +| **UXR-03**| Data Pipeline Shard Skipping / Re-reading. | Non-deterministic data loader state restoration upon checkpoint resume. | Silent loss spikes, training on duplicate batches, and data leakage. | **Suite C: DataLoader State Resume Suite** (Step-exact resume verification across hardware scales). | +| **UXR-04**| Numerical Divergence & FP8 Overflows. | Quantization scale underflows or overflows under mixed-precision formats. | Silent numerical corruption or `NaN` loss values at late training steps. | **Suite D: Precision Scale Boundary Suite** (Stress tests for quantization scale updates). | +| **UXR-05**| Unpredictable HBM Compilation OOMs. | Compilation memory footprint exceeds physical hardware capacity. | compilation OOM errors after spending 30+ minutes in XLA. | **Suite A: AOT Memory Profiler** (Static calculation of activation and weight footprints). | +| **UXR-06**| Distributed GRPO networking/connection dropouts. | Brittle connection/handshake management in high-rate multi-node training. | Intermittent connection errors that crash long-running RL/GRPO training workloads without retry logic. | **Suite E: Distributed Connection Resilience Test** (Stress-tests network drops and asserts auto-reconnection). | + +--- + +## 2. Proposed Specialized Corner & Integration Test Suites + +To address these findings, we propose seven targeted test suites integrated with standard pytest and CI/CD pipelines. + +```mermaid +graph TD + A[UXR Gaps & Ecosystem Friction] --> B[Corner & Integration Test Suites] + B --> C[Suite A: Config & Topology AOT] + B --> D[Suite B: Multi-Topology Resharding] + B --> E[Suite C: DataLoader Determinism] + B --> F[Suite D: Precision & Scale Stress] + B --> G[Suite E: Cross-Product Compatibility] + B --> H[Suite F: Docker, Logging & IAM] + B --> I[Suite G: Doc & CLI Linting] +``` + +### Suite A: Config AOT Fail-Fast & Topology Validator +Runs locally on CPU resources without requiring physical TPU hardware allocations. +* **Target Validation:** + * Asserts clean divisibility of dimension sizes (`base_emb_dim`, `base_mlp_dim`, `base_num_decoder_layers`) by computed mesh topology axes (`ici_tensor_parallelism`, `ici_fsdp_parallelism`). + * Calculates parameter counts and activation footprints statically to verify peak High Bandwidth Memory (HBM) against target hardware limits (e.g., TPU v5e/v6e, H100/H200). + * Rejects incompatible quantization and attention combinations (e.g., `quantization=int8` with custom attention kernels lacking INT8 support) during initial configuration parse. + +### Suite B: Multi-Topology Checkpoint Resharding Suite +Uses JAX's `jax.sharding` abstraction to simulate saving and loading checkpoints across virtual hardware layouts. +* **Target Validation:** + * **M-to-N scaling:** Simulates saving a checkpoint under FSDP topology (e.g., `fsdp_parallelism=8, tensor_parallelism=1`) and restoring it under an altered topology (e.g., `fsdp_parallelism=4, tensor_parallelism=2`). + * **Stack/Unstack checks:** Validates checkpoint serialization when moving between scanned (stacked) layers and unstacked layers for downstream tasks. + * **Preemption Recovery:** Asserts optimizer state (e.g., AdamW moment variables) and data markers restore identically from local emergency backups following simulated preemption signals. + +### Suite C: DataLoader Determinism & Shard Resume Suite +Verifies checkpoint restoration fidelity across HF and Grain data loader pipelines. +* **Target Validation:** + * **Step-Continuation Assert:** Asserts that running training continuously for 10 steps yields identical input tensors and loss values as running for 5 steps, checkpointing, resuming, and running another 5 steps. + * **Shard Boundary Assert:** Verifies that multiple host processes read non-overlapping data shards, and that resumption starts exactly at the global token index where the checkpoint was saved. + * **Empty Shard Handlers:** Simulates EOF boundary conditions and small shard sizes to ensure clean termination instead of process hangs. + +### Suite D: Precision Scale & Quantization Boundary Suite +Protects against numerical instability under low-precision configurations. +* **Target Validation:** + * **Extreme Value Probing:** Injects extreme activation values and verifies that scaling mechanisms (`absmax` or custom delayed updates) adjust dynamically without underflowing to zero or throwing `NaN`. + * **Quantization Transition:** Asserts correct dynamic application of INT8/FP8 layers when fine-tuning from non-quantized (e.g., BF16) checkpoints. + +### Suite E: Integrated Cross-Product Compatibility Suite +Implements the automated Pre-Release Signal pipeline to validate MaxText against its underlying ecosystem. +* **Target Validation:** + * **Pre-Release Hook Testing:** Triggers tests automatically when Release Candidates (RC) are created for core dependencies (vLLM, JAX, Tunix, Pathways, xpk). + * **Version Pinning validation:** Programmatically pins the RC version of the dependency in a sandboxed environment while retaining stable versions of other packages to isolate regressions. + * **Tutorial & Notebook Verification:** Automatically executes standard tutorials (Pre-Training, SFT, RL/GRPO, and Inference) to verify execution end-to-end. + * **Stacking Package Validation:** Audits environment setup scripts to verify that installing post-training libraries cleanly overrides existing pre-training packages (e.g., upgrading `tunix` cleanly without leaving conflicting legacy instances). + * **Notebook Profile Compatibility:** Automatically probes notebook hardware layouts to verify that the attached TPU devices meet the run prerequisites (e.g., verifying v5p availability). + * **Connection Resiliency Verification:** Simulates distributed network dropouts and jitter in multi-node training contexts (such as GRPO) to confirm auto-reconnect routines function correctly. + +### Suite F: Docker Build, Logging & IAM Permissions Diagnostic Suite +Minimizes setup delays, hidden failures, and permission blocks during local or remote container execution. +* **Target Validation:** + * **Lightweight SFT Builds:** Verifies that SFT builds utilize modular Dockerfile paths that exclude heavy, unnecessary components (e.g., `vLLM` or `tpu-common` when not in use) to reduce build times. + * **Runner Output Verbosity:** Verifies that Docker upload scripts and execution runners emit clear log levels, standard exit codes, and do not execute opaquely. + * **Pre-Execution IAM Probe:** Attempts a test write to the Artifact Registry using local credentials before invoking multi-step builds, ensuring users do not experience authorization failures at the end of the build cycle. + * **Pre-Execution Auth & Service Account Audits:** Verifies that local Docker authentication settings are correctly initialized (`gcloud auth configure-docker`) and that the VM-attached service account holds the proper permissions to write checkpoints to GCS and push to the registry. + * **State Persistence Assert:** Validates that XPK files and configurations remain persistent across simulated VM restarts. + * **Base VM Environment Sanity Checks:** Confirms that critical terminal tools and SSH access paths exist and operate under normal environments. + +### Suite G: Documentation, Configuration & Copy-Paste Validator +Validates tutorials, code templates, and CLI configurations to prevent syntax-level friction. +* **Target Validation:** + * **Tutorial Config Mapping:** Verifies that model configuration variables specified in documentation tutorials match the requirements of the accompanying code templates (e.g., preventing the use of non-instruct weights in instruct templates). + * **CLI Fallback Verification:** Asserts that CLI wrappers like XPK fall back to global environment configs (e.g., default `gcloud` zone) when command-specific flags are omitted. + * **Character Cleansing Test:** Automatically audits public documentation scripts to detect and strip hidden formatting or invalid newline characters from copy-paste containers. + * **Nomenclature Check:** Evaluates mapping files that translate TPU terminology between Cloud Console formats and GKE cluster definitions. + * **Prerequisite Path and Sequence Linting:** Confirms all prerequisite installations (e.g., cloning MaxText, configuring `uv` or dependencies) are ordered sequentially, with clear warnings on mutually exclusive options. + * **Post-Training Tutorial Integration:** Validates that architectural maps for downstream post-training techniques (GRPO, Distillation, PPO) are integrated without divergence. + +--- + +## 3. Dependency Matrix & Integration Topology + +To guide cross-product testing, the integration pipeline utilizes the following product dependency matrix: + +| Product | Depends On MaxText | Depends On Tunix | Depends On vLLM | Depends On tpu-inference | Depends On xpk | Depends On Pathways | Depends On JAX | +| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | +| **MaxText** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | +| **Tunix** | No | Yes | No | No | No | No | No | +| **vLLM** | No | No | Yes | Yes | Yes | Yes | Yes | +| **tpu-inference** | No | No | Yes | Yes | Yes | Yes | Yes | +| **xpk** | No | No | No | No | Yes | No | Yes | +| **Pathways** | No | No | No | No | No | Yes | Yes | +| **JAX** | No | No | No | No | No | No | Yes | + +--- + +## 4. Phased Implementation Plan + +We propose a four-phase integration roadmap to transition from localized validations to an ecosystem-wide automated CI pipeline: + +### Phase 1: Local AOT Validation & Tutorial Auditing (Weeks 1–2) +* **Action Items:** + * Implement structural mesh dimension checks in `src/maxtext/configs/pyconfig.py` to throw validation errors prior to JAX initialization. + * Integrate standard markdown validators in the documentation build to parse environment variable syntax, check for instruct/non-instruct config consistency, and verify chronological installation steps (e.g., ensuring tools like `uv` are not referenced prior to their setup step). + * Implement path checking in script runners to prevent execution failures when users trigger scripts from unexpected current working directories. + * Expose a command-line dry-run mode: `python3 src/maxtext/train.py base.yml dry_run=True` to compile the logical HLO graph on virtual devices and calculate expected memory. +* **Success Metric:** 100% of incompatible mesh, tutorial path, and prerequisite layout errors fail under 1 second on CPU, before container building or hardware compilation begins. + +### Phase 2: Checkpoint Resharding & Environment Stacking CI (Weeks 3–4) +* **Action Items:** + * Create automated pytest configurations that utilize virtual device grids to simulate resharding layouts. + * Implement dependency isolation checks within setup scripts to verify that installing new packages overrides legacy libraries (`tunix`) without package collision. + * Add pre-execution Docker credentials and service account capability audits to prevent late-stage image push blockages. + * Add standard pre-execution IAM checks in upload scripts to validate registry permission sets. +* **Success Metric:** Checkpoint restoration failures, container push auth blocks, and installation conflicts are caught in sandboxed CPU/virtual meshes prior to hardware allocation. + +### Phase 3: Modular Containers & Verbose Diagnostics (Weeks 5–6) +* **Action Items:** + * Refactor Dockerfile definitions to expose lightweight pathways for SFT and inference setups, eliminating unused package overhead. + * Incorporate standard logging layers, explicit exit codes, and verbose flags into the Docker upload runner utilities, providing time estimates for compilation/build milestones. + * Integrate data-loader step-resume validations and multi-node GRPO connection drop resilience tests into regular integration testing. + * Validate notebook setups against standard TPU VM environments to guarantee that linked templates match hardware profiles. +* **Success Metric:** SFT container compilation times are reduced by at least 2 minutes, and execution failures emit structured logs, progress metrics, and connection retry assertions. + +### Phase 4: Integrated Ecosystem CI & Compatibility Matrix (Weeks 7–8) +* **Action Items:** + * Deploy pre-release hooks that trigger automated testing when dependencies (JAX, vLLM, Tunix) publish Release Candidates. + * Configure a live, auto-updated compatibility dashboard showing verified versions: `MaxText vX.X is compatible with vLLM vY.Y, JAX vZ.Z, and Tunix vA.A`. +* **Success Metric:** 100% of breaking changes introduced by upstream dependencies are caught and notified before the dependency is released to the public. + +--- + +## 5. User Onboarding & First-Run Corner Cases + +When a developer first clones the MaxText repository, they are highly susceptible to environment-related failure modes. The test plan validates the following onboarding corner cases: + +### A. GCS Bucket Write & Permission Mismatches +* **The Failure:** Training compiles, but crashes hours later when attempting the first checkpoint because the environment's Application Default Credentials (ADC) lack write permissions to `base_output_directory`. +* **Validation:** Attempt to write a small, temporary dummy metadata file to the output directory during early startup to verify read/write permissions before graph compilation. + +### B. Gated Hugging Face Repository Credentials +* **The Failure:** Loading gated models (e.g., Llama, Gemma) fails late during execution because the user lacks an active `HF_TOKEN` or hasn't logged in via the Hugging Face CLI. +* **Validation:** Perform an early Hugging Face API request to assert credential validation and access permissions for the selected model family prior to initializing JAX layers. + +### C. Local Storage & Scratch Space Capacity +* **The Failure:** Local runs exhaust disk space when saving large checkpoints, causing process freezes or corrupting existing state. +* **Validation:** Calculate the expected size of checkpoints statically using parameter and optimizer dimensions. Ensure the local target disk has at least `3x` the target checkpoint size in free storage space. + +### D. Physical vs. Configured Device Mesh Mismatches +* **The Failure:** Executing workloads with a configured mesh size that does not align with the physical VM layout results in driver deadlocks. +* **Validation:** Prior to JAX compilation, verify that `prod(ici_fsdp_parallelism, ici_tensor_parallelism, ici_context_parallelism) == jax.local_device_count()`, throwing an immediate topology mismatch exception if they differ. + +### E. Dynamic Data Padding & Recompilation Loops +* **The Failure:** Shifting sequence lengths in input datasets trigger slow, repeated recompilations on every step. +* **Validation:** Monitor initial batch shapes. If sequence sizes vary across the first few steps, throw a diagnostic warning recommending padding configurations or `packing=True`. + +### F. Legacy Package Stacking Collisions +* **The Failure:** Upgraded tutorials or workflows crash due to legacy library versions (`tunix`) lingering in the environment. +* **Validation:** Add a pre-run script check that verifies system packages against active requirements, checking that older library instances are not shadowing the current package installations. + +### G. Stateful Tooling Persistence +* **The Failure:** CLI configurations or temporary elements established for cluster job execution (`xpk`) are cleared out after standard VM restarts. +* **Validation:** Implement persistence verification for configuration caches and tool assets, providing warning diagnostics if essential tools require re-initialization after a VM restart. + +### H. Pre-Flight Docker Authentication Verification +* **The Failure:** Docker image uploads fail with opaque permissions errors at the very end of a long image build because the docker daemon lacks setup credentials. +* **Validation:** Check for active registry endpoints in local Docker config files or run a dry-run credential check before launching build commands. + +### I. Service Account Permissions and Resource Scopes +* **The Failure:** VM workflows crash when retrieving files from GCS or writing images to the Artifact Registry because the attached VM service account lacks required scopes. +* **Validation:** Interrogate local metadata server endpoints during initial setup to assert that the active VM service account possesses correct IAM roles. + +### J. XPK Cluster and Workload Variable Pre-population +* **The Failure:** Submitting cluster tasks through XPK fails because expected variables (TPU type, workload ID, region parameters) are unset, forcing manual cloud console searches. +* **Validation:** Implement an onboarding dry-run script that probes local env variables and suggests export command overrides. + +### K. Notebook Hardware Profiler Sanity Check +* **The Failure:** Users running Colab notebooks hit driver or allocation errors due to the host executing on standard CPU/GPU instances instead of target TPU versions. +* **Validation:** Execute early, non-blocking cell checks to inspect `jax.devices()` and verify that physical accelerator shapes match expected TPU configs (e.g., v5p). + +### L. Distributed Connection Handshakes and Retries +* **The Failure:** Distributed workflows (such as GRPO) encounter transient network dropouts, immediately crashing hours of multi-host execution. +* **Validation:** Configure connection handshake retry budgets and keep-alive settings inside JAX/distributed runtime wrappers to gracefully survive transient drops. + +### M. Standard SSH & VM Package Manager Availability +* **The Failure:** Attempting to run automated setups or SSH overrides fails on stripped custom VM images because package managers (like `apt`) or standard connection paths are absent. +* **Validation:** Pre-flight shell capability checks to confirm minimal command sets are available, providing explicit, manual alternate recipes if executing inside constrained environments. + +--- + +> [!IMPORTANT] +> Integrating AOT validators and cross-dependency hooks directly into the local development and release workflows helps prevent the majority of runtime startup and update-related failures. + +> [!TIP] +> Using JAX's virtual device configurations, these suites can be evaluated within standard CPU-only environments, ensuring high execution velocity and low execution cost. diff --git a/resolve_conflict.py b/resolve_conflict.py new file mode 100644 index 0000000000..dd2f53c632 --- /dev/null +++ b/resolve_conflict.py @@ -0,0 +1,23 @@ +import sys + +def resolve(): + with open('src/maxtext/layers/quantizations.py', 'r') as f: + lines = f.readlines() + + out = [] + in_conflict = False + for line in lines: + if line.startswith('<<<<<<< HEAD'): + in_conflict = True + continue + if line.startswith('>>>>>>> origin/main'): + in_conflict = False + continue + + if not in_conflict: + out.append(line) + + with open('src/maxtext/layers/quantizations.py', 'w') as f: + f.writelines(out) + +resolve() diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 76f84e6ed8..b1d3e6bb15 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -393,9 +393,7 @@ def _load_params_nnx(self, params, rng): # axis metadata but no physical .sharding. Resolve logical to physical here so # device_put actually reshards instead of being a no-op. with nn_partitioning.axis_rules(self.config.logical_axis_rules): - target_shardings = sharding.nnx_construct_named_sharding( - params_abs, self._mesh - ) + target_shardings = sharding.nnx_construct_named_sharding(params_abs, self._mesh) params_state = jax.device_put(params, target_shardings) # We only need a concrete `rest` (RNG vars) for nnx.merge. create_nnx_sharded_model # builds the model with a jitted out_shardings so params are produced already @@ -403,9 +401,7 @@ def _load_params_nnx(self, params, rng): # large models). self.model is abstract with no .sharding, so pass an explicit one. _, full_abs = nnx.split(self.model) with nn_partitioning.axis_rules(self.config.logical_axis_rules): - full_sharding = sharding.nnx_construct_named_sharding( - full_abs, self._mesh - ) + full_sharding = sharding.nnx_construct_named_sharding(full_abs, self._mesh) concrete_model = maxtext_utils_nnx.create_nnx_sharded_model( self.model, self._create_model_fn, mesh=self._mesh, named_sharding=full_sharding ) diff --git a/src/maxtext/layers/nnx_wrappers.py b/src/maxtext/layers/nnx_wrappers.py index 3d17ce0ea4..c08011c062 100644 --- a/src/maxtext/layers/nnx_wrappers.py +++ b/src/maxtext/layers/nnx_wrappers.py @@ -555,8 +555,10 @@ def __getattr__(self, name: str): return self.kwargs[name] maybe_method = getattr(self.nnx_class, name, None) if callable(maybe_method): + def unbound_func(instance, *args, **kwargs): return instance.__call__(*args, nnx_method=maybe_method, **kwargs) + return types.MethodType(unbound_func, self) return super().__getattribute__(name) @@ -574,7 +576,6 @@ def _update_variables(self, module): collection_flat_state[collection] = [] collection_flat_state[collection].append((path, leaf)) - # update linen variables for collection, flat_state in collection_flat_state.items(): if self.is_mutable_collection(collection): diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 8aa9c95659..a8cc822758 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -1289,7 +1289,7 @@ def _traverse(current_state, current_path): for k, v in list(current_state.items()): new_path = current_path + (k,) if isinstance(v, nnx.Intermediate): - if len(new_path) >= len(suffix) and new_path[-len(suffix):] == suffix: + if len(new_path) >= len(suffix) and new_path[-len(suffix) :] == suffix: results.append((new_path, v)) if clear: keys_to_delete.append(k) @@ -1346,12 +1346,14 @@ def get_intermediate_value(model, nested_key, default=None, clear=False): # Helper key function to sort paths numerically if indices are present def path_sort_key(item): path = item[0] + def _to_int_if_possible(val): if isinstance(val, int): return val if isinstance(val, str) and val.isdigit(): return int(val) return val + return tuple(_to_int_if_possible(x) for x in path) found.sort(key=path_sort_key) @@ -1710,18 +1712,14 @@ def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=Tru # ourselves via nnx_construct_named_sharding, so auto-assignment is not needed here. abs_model = nnx.eval_shape(nnx_init_trainstate_fn) _, abs_var_state = nnx.split(abs_model) - named_sharding_state = sharding.nnx_construct_named_sharding( - abs_var_state, mesh - ) + named_sharding_state = sharding.nnx_construct_named_sharding(abs_var_state, mesh) abstract_state = jax.tree.map( lambda a, s: jax.ShapeDtypeStruct(a.shape, a.dtype, sharding=s), abs_var_state, named_sharding_state, ) - state_mesh_shardings = maxtext_utils_nnx.nnx_extract_named_sharding( - abstract_state - ) + state_mesh_shardings = maxtext_utils_nnx.nnx_extract_named_sharding(abstract_state) if is_training and config.shard_optimizer_over_data: # Add data to sharding for optimizer state diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 6449c890b9..be2c3c4e47 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -19,7 +19,7 @@ from typing import Any, Sequence import unittest import pytest -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock, patch from dataclasses import dataclass, field import numpy as np import optax @@ -95,15 +95,22 @@ def test_fp8_stats_not_clipped_but_others_are(self): ) ) + class MockAttention(nnx.Module): + """Mock attention module for testing intermediate extraction.""" + def __init__(self, rngs): pass + def __call__(self, x, sow_val=None): if sow_val is not None: self.sow(nnx.Intermediate, "out_projection_activations", sow_val) return x + class MockDecoderLayer(nnx.Module): + """Mock decoder layer for testing intermediate extraction.""" + def __init__(self, rngs, attention_name="self_attention"): if attention_name == "self_attention": self.self_attention = MockAttention(rngs) @@ -115,12 +122,17 @@ def __call__(self, x, sow_val=None): attention = getattr(self, self.attention_name) return attention(x, sow_val) + class MockDecoderSequential(nnx.Module): + """Mock decoder sequential block for testing intermediate extraction.""" + def __init__(self, rngs, attention_name="self_attention"): - self.layers = nnx.List([ - MockDecoderLayer(rngs, attention_name), - MockDecoderLayer(rngs, attention_name), - ]) + self.layers = nnx.List( + [ + MockDecoderLayer(rngs, attention_name), + MockDecoderLayer(rngs, attention_name), + ] + ) def __call__(self, x, sow_vals=None): if sow_vals is None: @@ -130,7 +142,10 @@ def __call__(self, x, sow_vals=None): out = layer(out, val) return out + class MockDecoderScanned(nnx.Module): + """Mock decoder scanned block for testing intermediate extraction.""" + def __init__(self, rngs, attention_name="self_attention"): self.layers = MockDecoderLayer(rngs, attention_name) self.attention_name = attention_name @@ -141,7 +156,10 @@ def __call__(self, x, sow_val=None): attention.sow(nnx.Intermediate, "out_projection_activations", sow_val) return x + class MockTransformer(nnx.Module): + """Mock transformer for testing intermediate extraction.""" + def __init__(self, decoder_type, rngs, attention_name="self_attention"): if decoder_type == "sequential": self.decoder = MockDecoderSequential(rngs, attention_name) @@ -1600,9 +1618,7 @@ class MockConfig: optimizer_memory_host_offload: bool = False parameter_memory_host_offload: bool = False param_scan_axis: int = 0 - logical_axis_rules: list = field( - default_factory=lambda: [["data", ["data"]], ["model", ["model"]]] - ) + logical_axis_rules: list = field(default_factory=lambda: [["data", ["data"]], ["model", ["model"]]]) class MockTrainState(nnx.Module): """Simulates a TrainState with params and optimizer state.""" diff --git a/tests/unit/tiling_test.py b/tests/unit/tiling_test.py index 0473e797e8..a302a21a65 100644 --- a/tests/unit/tiling_test.py +++ b/tests/unit/tiling_test.py @@ -337,7 +337,6 @@ def test_vocab_tiling_gradient_with_z_loss_nnx(self): "Gradients do not match for vocab tiling when z-loss is enabled (NNX).", ) - @pytest.mark.tpu_only def test_vocab_tiling_nnx_loss(self): """ From b9fca6a4bed37a29db40215a5228b1749e7c8839 Mon Sep 17 00:00:00 2001 From: Sarun Singla Date: Thu, 25 Jun 2026 15:57:21 +0000 Subject: [PATCH 52/52] Remove accidentally committed scratch files --- .../aqt_deprecation_phase_2_execution_plan.md | 187 -------------- docs/uxr_usability_test_plan.md | 235 ------------------ resolve_conflict.py | 23 -- 3 files changed, 445 deletions(-) delete mode 100644 docs/aqt_deprecation_phase_2_execution_plan.md delete mode 100644 docs/uxr_usability_test_plan.md delete mode 100644 resolve_conflict.py diff --git a/docs/aqt_deprecation_phase_2_execution_plan.md b/docs/aqt_deprecation_phase_2_execution_plan.md deleted file mode 100644 index 86be007916..0000000000 --- a/docs/aqt_deprecation_phase_2_execution_plan.md +++ /dev/null @@ -1,187 +0,0 @@ -# MaxText AQT Deprecation Phase #2: Technical Execution Plan - -This document provides a comprehensive, highly technical, and definitive execution plan for Phase #2 of the Accurate Quantization Training (AQT) deprecation within the MaxText repository. - -Our primary goal is to **completely strip out the legacy AQT-specific code, configurations, and dependencies** while **maintaining absolute architectural generality** for modern quantization backends (such as Qwix, native FP8, and TransformerEngine). - ---- - -## 1. Architectural Code Review & Impact Assessment - -We have performed a deep-dive architectural review of the AQT footprint across the MaxText repository. Below is the precise mapping of AQT dependencies, configuration flags, and the "blast radius" of their removal. - -```mermaid -graph TD - subgraph Configuration - types_py["configs/types.py"] --> |Defines| QuantType["QuantizationType Enum"] - types_py --> |Defines| QuantConfig["Quantization Class"] - end - - subgraph Core Quantization - quant_py["layers/quantizations.py"] --> |Imports| AQT_v2["aqt.jax.v2"] - quant_py --> |Implements| AqtQuant["AqtQuantization"] - quant_py --> |Implements| MemoryUtils["remove_quantized_params"] - end - - subgraph Layers & Models - linears_py["layers/linears.py"] --> |Calls| AqtQuant - models_py["models/models.py"] --> |Type Hints| AqtQuant - decoders_py["layers/decoders.py"] --> |Type Hints| AqtQuant - end - - subgraph Tooling & Inference - lw_quant["utils/layerwise_quantization.py"] --> |Imports| AQT_v2 - maxengine["inference/maxengine/maxengine.py"] --> |Calls| MemoryUtils - end - - AqtQuant -.-> |Injects| linears_py - AqtQuant -.-> |Injects| models_py - AqtQuant -.-> |Injects| decoders_py -``` - -### A. Specific Modules, Files, and Configuration Flags to Modify/Remove - -1. **External Library Dependencies:** - * **File:** [requirements.txt](file:///usr/local/google/home/sarunsingla/maxtext/src/dependencies/requirements/requirements.txt#L2), [base_requirements/requirements.txt](file:///usr/local/google/home/sarunsingla/maxtext/src/dependencies/requirements/base_requirements/requirements.txt#L2), and generated cuda/tpu lockfiles. - * **Action:** Remove the `aqtp` package dependency (which installs `aqt` in Python). -2. **Configuration Flags & Enums:** - * **File:** [types.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py) - * **Flags to Remove:** - * `replicate_quant_scale`: AQT-specific flag used to replicate scales across mesh axes to avoid inefficient XLA fusions. - * `quant_cfg_path`: Path for `intmp` (AQT mixed precision) configurations. - * **Quantization Types to Prune:** Remove `INTMP = "intmp"`, `aqt_fp8`, and `aqt_fp8_full` from [QuantizationType](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py#L83) enum and config parsing. - * **Flag to Deprecate/Modify:** - * `use_qwix_quantization`: Currently defaults to `False`. For Phase 2, we will set this default to `True` or completely remove it, making Qwix/native FP8 the standard quantization pipeline. -3. **Core Quantization Library:** - * **File:** [quantizations.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py) - * **Action:** - * Remove all `aqt.jax.v2` and `aqt_flax` imports. - * Delete the [AqtQuantization](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py#L119) dataclass. - * Delete AQT helper functions: `_tiling_fn`, `_rhs_axis_metadata_wrapper`, `_build_const_scale_config`, `_build_per_tensor_config`, `_get_int8_quant_config`, `_get_aqt_fp8_quant_config`, `_get_aqt_fp8_default_config`, `_dot_general_make`, `_get_default_mp_config`, `_get_mixed_precision_quant_config`. - * Delete AQT memory optimization utilities: `match_aqt_and_unquantized_param`, `_get_aqt_key_paths`, and [remove_quantized_params](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py#L695). -4. **Inference Engine Integration:** - * **File:** [maxengine.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/inference/maxengine/maxengine.py#L524) - * **Action:** Remove the call to `quantizations.remove_quantized_params` and the handling of the `"aqt"` variable collection. -5. **Layer-wise Quantization Tooling:** - * **File:** [layerwise_quantization.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/utils/layerwise_quantization.py) - * **Action:** This utility is deeply coupled to AQT's `QTensor` and `AqtDotGeneral` structures. We will completely deprecate this file or refactor it to target Qwix/native FP8. - -### B. Blast Radiuses & Critical Decoupling -* **The Quantization Param Type Hint:** Over 20 files (including [models.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/models/models.py#L36), [linears.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/linears.py#L37), [decoders.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/decoders.py#L40), and [moe.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/moe.py#L192)) import and type-hint the `quant` parameter using `AqtQuantization as Quant`. - * *Mitigation:* We must **not** delete the base class [Quantization](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py#L55). Instead, we will import and type-hint with `Quantization as Quant`. This keeps layer signatures clean and compatible with Qwix, native FP8, and TransformerEngine. -* **Layer Computations:** In [linears.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/linears.py#L70), `_compute_dot_general` and `_compute_dot_general_nnx` invoke `quant.dot_general_cls()`. - * *Mitigation:* As long as other quantization backends inherit from the base `Quantization` class and implement `dot_general_cls()`, their execution remains completely unaffected. - ---- - -## 2. Phase #2 Deprecation Plan & Scope of Changes - -We will execute the AQT removal using a structured, phased engineering sprint. - -### Step 1: Remove Third-Party Dependency & Clean Configs -1. Prune `aqtp` from all requirements lockfiles. -2. Edit [types.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py) to remove `replicate_quant_scale` and `quant_cfg_path` from the [Quantization](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py#L423) configuration class. -3. Remove the AQT warning check from `set_derived_and_validate_values` in [types.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/configs/types.py#L2574-L2580). -4. Throw a compile-time `ValueError` if a user attempts to run quantization with `use_qwix_quantization=False` and the quantization type is not native FP8 or TransformerEngine. - -### Step 2: Refactor `layers/quantizations.py` -Modify [quantizations.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/quantizations.py) to remove all AQT references: - -```diff --from aqt.jax.v2 import config as aqt_config --from aqt.jax.v2 import aqt_tensor --from aqt.jax.v2.flax import aqt_flax --from aqt.jax.v2 import tiled_dot_general --from aqt.jax.v2 import calibration - -... - --@dataclass --class AqtQuantization: -- """Configures AQT quantization github.com/google/aqt.""" -- ... -- def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): -- ... -- def einsum(self, mesh_axes: Tuple[str, ...] = ()): -- ... - -... - --def remove_quantized_params(params, aqt_vars): -- ... -``` - -### Step 3: Generalize Layer and Model Type-Hints -In every model and layer file, update the imports and type hints: - -```diff --from maxtext.layers.quantizations import AqtQuantization as Quant -+from maxtext.layers.quantizations import Quantization as Quant -``` - -This change preserves the signature of all layers: -```python -class DenseGeneral(nnx.Module): - def __init__( - self, - ... - quant: None | Quant = None, - ... - ): -``` - -### Step 4: Prune Inference & Tooling Hooks -1. In [maxengine.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/inference/maxengine/maxengine.py#L524), remove AQT parameter extraction and the `remove_quantized_params` call. -2. Delete or refactor [layerwise_quantization.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/utils/layerwise_quantization.py) to utilize Qwix quantization rules instead of AQT's `QTensor` structures. - ---- - -## 3. End-to-End (E2E) Testing Strategy - -To guarantee that this cleanup introduces zero regressions in model correctness, hardware utilization, or training stability, we will execute the following multi-tier verification matrix. - -``` - [Unit Tests] [Integration Tests] [E2E Training Runs] - (quantizations_test.py) (maxengine_test.py) (Llama2-7B / Gemma2-9B) - │ │ │ - ├─► Prune AQT tests ├─► Test Qwix/FP8 ├─► Verify Convergence - └─► Verify Qwix/FP8 └─► Verify serving latency └─► Check TPU/GPU MFU -``` - -### A. Unit & Integration Tests -* **File to Refactor:** [quantizations_test.py](file:///usr/local/google/home/sarunsingla/maxtext/tests/unit/quantizations_test.py) - * **Action:** - * Remove `QuantTestModule`'s explicit dependency on `aqt_flax.AqtDotGeneral` and `aqt_flax.AqtEinsum`. - * Prune AQT-specific test cases: `test_aqt_quantization`, `test_mixed_precision_*`, and `test_remove_quantized_params`. - * Enhance existing Qwix and native FP8 tests (`test_int8_quantization`, `test_fp8_quantization`, `test_fp8_full_quantization`) to ensure full coverage of the remaining quantization pathways. -* **Other Unit Tests:** Run `tiling_test.py`, `gpt3_test.py`, `maxtext_utils_test.py`, and `model_test.py` to ensure that standard unquantized runs are completely unaffected. - -### B. Regression Testing (Performance & Correctness) -* **Convergence Verification:** Run a Llama2-7B training run for 1000 steps using BF16 (unquantized) and Qwix FP8. Compare the loss curves before and after the PR to verify mathematical equivalence. -* **Throughput & Utilization:** Monitor Step Time (ms) and Model Flops Utilization (MFU) on TPU v5e/v6e and NVIDIA H100 GPUs. The elimination of AQT code paths must result in **zero throughput regression** and **zero HBM memory overhead increases** for standard models. - -### C. E2E Training Scale Verification -We will run E2E training workloads to validate the entire JIT compilation and execution pipeline: -1. **Workload A (TPU Scale):** Llama2-7B training on a 16-device TPU v6e slice using Qwix FP8 (`quantization=fp8_full` and `use_qwix_quantization=True`). -2. **Workload B (GPU Scale):** Gemma2-9B training on an 8-GPU H100 node using native FP8 (`quantization=fp8` and `use_qwix_quantization=True`). -3. **Workload C (MoE Scale):** DeepSeek-V3 MoE model initialized and executed using TransformerEngine quantization (`quantization=te_fp8_delayedscaling`). - ---- - -## 4. Anticipated Edge Cases, Bug Fixing, and Risk Mitigation - -During this infrastructure migration, we anticipate 3 major technical hurdles. Below are the exact engineering remedies for each. - -### Hurdle 1: Legacy Checkpoint Loading Failure (Structure Mismatch) -* **The Issue:** Legacy Orbax checkpoints saved using AQT quantization contain AQT-specific variables (`qrhs` and `frozen` QTensor structures). Attempting to restore these checkpoints into the new AQT-free model structure will cause Orbax to throw a structural mismatch error because the model no longer defines those variables. -* **The Remedy:** - * Implement an explicit warning or error in [model_creation_utils.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/utils/model_creation_utils.py#L334) during the checkpoint restoration phase. - * Provide a standalone migration script (`maxtext/checkpoint_conversion/remove_aqt_vars.py`) that loads a legacy AQT checkpoint, strips out the AQT-specific variables (converting them back to unquantized parameters or Qwix-compatible weights), and saves a clean Orbax checkpoint. - -### Hurdle 2: Shape Mismatches in Grouped MatMul (GMM) / MoE Kernels -* **The Issue:** AQT deprecation could inadvertently break how tiling dimensions and contraction axes are set up in MoE layers, leading to shape mismatches in custom Pallas kernels (like GMM in [ops.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/kernels/megablox/ops.py)). -* **The Remedy:** Ensure that the generic `Quantization` interface correctly plumbs the tile sizes and contraction axes to `DenseGeneral` and `MlpBlock` in [linears.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/linears.py#L188). Add strict shape validation assertions in [moe.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/moe.py) to verify that the intermediate sharding matches the GMM kernel expectations before launching the kernel. - -### Hurdle 3: Compile-time Tracing Failures in NNX Bridge -* **The Issue:** The NNX-to-Linen bridge ([nnx_wrappers.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/nnx_wrappers.py)) traces the module during `lazy_init` and expects certain mutable collections (like `"aqt"`). If these collections are missing or empty, JAX might throw a tracing error. -* **The Remedy:** In [linears.py](file:///usr/local/google/home/sarunsingla/maxtext/src/maxtext/layers/linears.py#L97), update `_compute_dot_general_nnx` to only include `"aqt"` in the `mutable` list if the quantization backend actually registers that collection. For Qwix and native FP8, ensure they use their respective state collections (like `"params"` or `"intermediates"`), preventing JAX tracing crashes. diff --git a/docs/uxr_usability_test_plan.md b/docs/uxr_usability_test_plan.md deleted file mode 100644 index 91e03550cd..0000000000 --- a/docs/uxr_usability_test_plan.md +++ /dev/null @@ -1,235 +0,0 @@ -# UXR Usability Analysis & Ecosystem-Aligned Corner Test Plan for MaxText - -This document presents an analysis of user-facing pain points, usability gaps, and cross-product dependency friction modes in MaxText. It proposes a concrete, phased implementation plan for introducing **specialized corner test suites** and an **integrated cross-product CI pipeline** designed to prevent regressions, improve developer velocity, and ensure stable training and inference runs on large-scale TPU/GPU topologies. - ---- - -## 1. Integrated UXR & Ecosystem Friction Analysis - -MaxText operates within a multi-layered ecosystem of high-performance libraries for training, inference, orchestration, and job submission. Because dependencies release on independent cycles, uncoordinated updates frequently introduce regressions. - -The following table details usability gaps, environment friction points (MT-01 to MT-19), and the proposed testing/CI remedies to shift validation left: - -| ID | Friction / Usability Pain Point | Technical Root Cause & Dependencies | Diagnostic / Impact | Proposed Testing Remedy | -| :--- | :--- | :--- | :--- | :--- | -| **MT-01** | VM restarts wipe newly installed XPK elements. | State non-persistence across VM restarts. | Users lose environment setup, forcing rebuilds. | **Suite F: VM State Persistence Test** (Asserts configuration and tool persistence across simulated restarts). | -| **MT-02** | Stacking post-training packages fails to override older library versions. | Installation conflicts with pre-existing packages (specifically `tunix`). | Opaque Python keyword errors that halt post-training setups. | **Suite E: Environment Package Isolation Audit** (Validates clean overrides and environment resets). | -| **MT-03** | GRPO/SFT workflows require building full Docker images from source. | Absence of pre-built, modular release images and clear build progress. | 10+ minute delay per execution attempt with no progress feedback or build time estimates. | **Suite F: Modular Docker & Pre-built Release Validation** (Verifies compatibility of pre-built images and build timeline output). | -| **MT-04** | Unused heavy dependencies (vLLM, tpu-common) are installed during SFT. | Non-modular Dockerfile configurations. | Unnecessary image compilation overhead (adds 2+ minutes per build). | **Suite F: Lightweight Build Pathway Test** (Validates dependency footprint optimization in SFT Dockerfiles). | -| **MT-05** | Docker upload runner script fails silently. | Opaque utility runner execution without logs or stdout capture. | Users unable to self-diagnose; requires manual code edits to trace. | **Suite F: Verbose Log & Exit Code Auditor** (Enforces standard error logging and explicit exit codes). | -| **MT-06** | Documentation interchanges critical environment variables. | Inconsistent naming convention in tutorials (e.g., `MODEL` vs `MODEL_NAME`). | Continuous user copy-paste failures and runtime syntax errors. | **Suite G: Documentation Linting Suite** (Automated markdown command syntax and variable validation). | -| **MT-07** | Tutorial configurations specify non-instruct models with instruct templates. | Mismatched model weight configs and runtime chat templates. | Immediate execution crashes when running tutorials. | **Suite G: Model Configuration Auditor** (Validates matching characteristics between configs and code). | -| **MT-08** | Prerequisite installation steps are buried or out-of-order in explanatory prose. | Suboptimal layout structure and sequencing in user tutorials. | Users skim text, miss vital setup commands (e.g., installing `uv`), or try to run utilities before they are installed. | **Suite G: Document Structure Auditor** (Enforces separation and logical order of execution steps). | -| **MT-09** | XPK commands require redundant `--zone` flags. | CLI does not inherit global `gcloud` configurations. | Redundant input overhead and minor command-line syntax friction. | **Suite G: CLI Argument Inheritance Test** (Ensures CLI tools respect global fallback configurations). | -| **MT-10** | TPU cluster nomenclature is inconsistent across platforms. | Discrepant naming schemas (e.g., Console `v6e-8` vs GKE `ct6e-standard-4t`). | Hardware configuration confusion during cluster mapping. | **Suite G: Naming Mapping Validator** (Checks correctness of taxonomy mapping utilities). | -| **MT-11** | Copied multi-line code blocks inject hidden newline characters. | Format copying issues in public documentation hub. | Compound execution and string parsing errors in terminal sessions. | **Suite G: Copy-Paste Character Sanitizer** (Tests copy-paste text scrubbing on documentation assets). | -| **MT-12** | Ambiguity regarding git cloning and install choices. | Documentation fails to clarify when steps are mutually exclusive or mandatory. | Users perform redundant setup steps (e.g., installing both vLLM and TPU branches) or miss foundational repository checkouts. | **Suite G: Prerequisite Path Validator** (Verifies dependency branches are distinct and explicitly mutually exclusive). | -| **MT-13** | Missing critical setup templates and commands. | Documentation omits fundamental infrastructure preparation steps (e.g., GCS bucket creation for SFT output). | Users are blocked or forced to construct cloud setup commands manually from scratch. | **Suite G: Command Template Auditor** (Verifies documentation contains copy-pastable templates for all prerequisites). | -| **MT-14** | Path assumptions cause execution errors. | Scripts assume a specific working directory without validating the user's path context. | Users run commands from non-root paths and receive path errors (e.g., with `docker_upload_runner.sh`). | **Suite G: Directory-Agnostic Runner Test** (Validates that utility runners execute successfully regardless of invocation CWD). | -| **MT-15** | Context switching across fragmented documentation silos. | Information is scattered across MaxText GitHub, Cloud, and external libraries. | High cognitive load and slower setup times as users hop between distinct doc ecosystems. | **Suite G: Cross-Doc Link Integrity Suite** (Validates correctness and co-location of essential setup links). | -| **MT-16** | Divergent tutorials for expanding post-training techniques. | Lack of centralized architectural strategy for post-training guides (GRPO, Distillation, PPO, DPO). | Maintenance burden increases and documentation diverges as more techniques are added. | **Suite G: Post-Training Documentation Matrix** (Ensures unified branching structure for post-training guides). | -| **MT-17** | XPK commands lack clear env variable configuration guidance. | Setup scripts fail to document or auto-populate necessary cluster and workload variables. | Users must manually search Cloud Console to find details like TPU type or workload ID to fill XPK parameters. | **Suite F: XPK Environment Variable Audit** (Validates existence of helper env scripts and config parsers). | -| **MT-18** | Incompatible TPU profiles in notebooks. | Colab notebooks contain default links with unsupported or unavailable TPU types (e.g., requesting v5p without access). | Immediate failure during first cell execution in Colab, blocking users. | **Suite E: Notebook Hardware Compatibility Test** (Validates TPU availability matches notebook requirements). | -| **MT-19** | Package manager and SSH command failures on locked VMs. | Standard tools (`sudo apt`) are missing or standard SSH structures fail in stripped environments. | Users cannot update libraries or connect to VMs, leading to absolute blockages. | **Suite F: Base Image Tooling Sanity Test** (Checks that minimal required CLI commands are present and functional). | -| **SEC-01**| Project Editor roles lack write access to the Artifact Registry. | Missing IAM permission guidance during Docker image uploads. | Users blocked immediately on image push, forcing manual IAM fixes. | **Suite F: Pre-Execution IAM Validator** (Checks Registry Writer access before starting upload scripts). | -| **SEC-02**| Missing Docker authentication instructions. | Guide omits required configuration commands (e.g., `gcloud auth configure-docker`). | Silent permission denied errors at the final image upload step. | **Suite F: Pre-Execution Auth Validator** (Validates docker auth config prior to running build actions). | -| **SEC-03**| Opaque TPU VM service account mapping. | Documentation lacks clear guidelines on VM-attached service accounts vs personal credentials. | Standard workflows fail on permission checks, requiring manual troubleshooting. | **Suite F: Service Account IAM Auditor** (Verifies attached TPU VM service accounts possess the required GCS/Registry scopes). | -| **UXR-01**| Late-Stage Configuration Failures. | No ahead-of-time (AOT) configuration validation. | Wasted TPU resource hours; errors caught hours into compilation or step runs. | **Suite A: Static Config & Topology Validator** (AOT checks for mesh dimension alignment and memory). | -| **UXR-02**| Restoration Hangs & Resharding OOMs. | Mesh topology mismatches (e.g., 2D to 3D) during checkpoint restoration. | Silent compilation hangs or abrupt OOM errors during restoration. | **Suite B: Checkpoint Resharding Suite** (Mock-reshard testing across virtual topologies). | -| **UXR-03**| Data Pipeline Shard Skipping / Re-reading. | Non-deterministic data loader state restoration upon checkpoint resume. | Silent loss spikes, training on duplicate batches, and data leakage. | **Suite C: DataLoader State Resume Suite** (Step-exact resume verification across hardware scales). | -| **UXR-04**| Numerical Divergence & FP8 Overflows. | Quantization scale underflows or overflows under mixed-precision formats. | Silent numerical corruption or `NaN` loss values at late training steps. | **Suite D: Precision Scale Boundary Suite** (Stress tests for quantization scale updates). | -| **UXR-05**| Unpredictable HBM Compilation OOMs. | Compilation memory footprint exceeds physical hardware capacity. | compilation OOM errors after spending 30+ minutes in XLA. | **Suite A: AOT Memory Profiler** (Static calculation of activation and weight footprints). | -| **UXR-06**| Distributed GRPO networking/connection dropouts. | Brittle connection/handshake management in high-rate multi-node training. | Intermittent connection errors that crash long-running RL/GRPO training workloads without retry logic. | **Suite E: Distributed Connection Resilience Test** (Stress-tests network drops and asserts auto-reconnection). | - ---- - -## 2. Proposed Specialized Corner & Integration Test Suites - -To address these findings, we propose seven targeted test suites integrated with standard pytest and CI/CD pipelines. - -```mermaid -graph TD - A[UXR Gaps & Ecosystem Friction] --> B[Corner & Integration Test Suites] - B --> C[Suite A: Config & Topology AOT] - B --> D[Suite B: Multi-Topology Resharding] - B --> E[Suite C: DataLoader Determinism] - B --> F[Suite D: Precision & Scale Stress] - B --> G[Suite E: Cross-Product Compatibility] - B --> H[Suite F: Docker, Logging & IAM] - B --> I[Suite G: Doc & CLI Linting] -``` - -### Suite A: Config AOT Fail-Fast & Topology Validator -Runs locally on CPU resources without requiring physical TPU hardware allocations. -* **Target Validation:** - * Asserts clean divisibility of dimension sizes (`base_emb_dim`, `base_mlp_dim`, `base_num_decoder_layers`) by computed mesh topology axes (`ici_tensor_parallelism`, `ici_fsdp_parallelism`). - * Calculates parameter counts and activation footprints statically to verify peak High Bandwidth Memory (HBM) against target hardware limits (e.g., TPU v5e/v6e, H100/H200). - * Rejects incompatible quantization and attention combinations (e.g., `quantization=int8` with custom attention kernels lacking INT8 support) during initial configuration parse. - -### Suite B: Multi-Topology Checkpoint Resharding Suite -Uses JAX's `jax.sharding` abstraction to simulate saving and loading checkpoints across virtual hardware layouts. -* **Target Validation:** - * **M-to-N scaling:** Simulates saving a checkpoint under FSDP topology (e.g., `fsdp_parallelism=8, tensor_parallelism=1`) and restoring it under an altered topology (e.g., `fsdp_parallelism=4, tensor_parallelism=2`). - * **Stack/Unstack checks:** Validates checkpoint serialization when moving between scanned (stacked) layers and unstacked layers for downstream tasks. - * **Preemption Recovery:** Asserts optimizer state (e.g., AdamW moment variables) and data markers restore identically from local emergency backups following simulated preemption signals. - -### Suite C: DataLoader Determinism & Shard Resume Suite -Verifies checkpoint restoration fidelity across HF and Grain data loader pipelines. -* **Target Validation:** - * **Step-Continuation Assert:** Asserts that running training continuously for 10 steps yields identical input tensors and loss values as running for 5 steps, checkpointing, resuming, and running another 5 steps. - * **Shard Boundary Assert:** Verifies that multiple host processes read non-overlapping data shards, and that resumption starts exactly at the global token index where the checkpoint was saved. - * **Empty Shard Handlers:** Simulates EOF boundary conditions and small shard sizes to ensure clean termination instead of process hangs. - -### Suite D: Precision Scale & Quantization Boundary Suite -Protects against numerical instability under low-precision configurations. -* **Target Validation:** - * **Extreme Value Probing:** Injects extreme activation values and verifies that scaling mechanisms (`absmax` or custom delayed updates) adjust dynamically without underflowing to zero or throwing `NaN`. - * **Quantization Transition:** Asserts correct dynamic application of INT8/FP8 layers when fine-tuning from non-quantized (e.g., BF16) checkpoints. - -### Suite E: Integrated Cross-Product Compatibility Suite -Implements the automated Pre-Release Signal pipeline to validate MaxText against its underlying ecosystem. -* **Target Validation:** - * **Pre-Release Hook Testing:** Triggers tests automatically when Release Candidates (RC) are created for core dependencies (vLLM, JAX, Tunix, Pathways, xpk). - * **Version Pinning validation:** Programmatically pins the RC version of the dependency in a sandboxed environment while retaining stable versions of other packages to isolate regressions. - * **Tutorial & Notebook Verification:** Automatically executes standard tutorials (Pre-Training, SFT, RL/GRPO, and Inference) to verify execution end-to-end. - * **Stacking Package Validation:** Audits environment setup scripts to verify that installing post-training libraries cleanly overrides existing pre-training packages (e.g., upgrading `tunix` cleanly without leaving conflicting legacy instances). - * **Notebook Profile Compatibility:** Automatically probes notebook hardware layouts to verify that the attached TPU devices meet the run prerequisites (e.g., verifying v5p availability). - * **Connection Resiliency Verification:** Simulates distributed network dropouts and jitter in multi-node training contexts (such as GRPO) to confirm auto-reconnect routines function correctly. - -### Suite F: Docker Build, Logging & IAM Permissions Diagnostic Suite -Minimizes setup delays, hidden failures, and permission blocks during local or remote container execution. -* **Target Validation:** - * **Lightweight SFT Builds:** Verifies that SFT builds utilize modular Dockerfile paths that exclude heavy, unnecessary components (e.g., `vLLM` or `tpu-common` when not in use) to reduce build times. - * **Runner Output Verbosity:** Verifies that Docker upload scripts and execution runners emit clear log levels, standard exit codes, and do not execute opaquely. - * **Pre-Execution IAM Probe:** Attempts a test write to the Artifact Registry using local credentials before invoking multi-step builds, ensuring users do not experience authorization failures at the end of the build cycle. - * **Pre-Execution Auth & Service Account Audits:** Verifies that local Docker authentication settings are correctly initialized (`gcloud auth configure-docker`) and that the VM-attached service account holds the proper permissions to write checkpoints to GCS and push to the registry. - * **State Persistence Assert:** Validates that XPK files and configurations remain persistent across simulated VM restarts. - * **Base VM Environment Sanity Checks:** Confirms that critical terminal tools and SSH access paths exist and operate under normal environments. - -### Suite G: Documentation, Configuration & Copy-Paste Validator -Validates tutorials, code templates, and CLI configurations to prevent syntax-level friction. -* **Target Validation:** - * **Tutorial Config Mapping:** Verifies that model configuration variables specified in documentation tutorials match the requirements of the accompanying code templates (e.g., preventing the use of non-instruct weights in instruct templates). - * **CLI Fallback Verification:** Asserts that CLI wrappers like XPK fall back to global environment configs (e.g., default `gcloud` zone) when command-specific flags are omitted. - * **Character Cleansing Test:** Automatically audits public documentation scripts to detect and strip hidden formatting or invalid newline characters from copy-paste containers. - * **Nomenclature Check:** Evaluates mapping files that translate TPU terminology between Cloud Console formats and GKE cluster definitions. - * **Prerequisite Path and Sequence Linting:** Confirms all prerequisite installations (e.g., cloning MaxText, configuring `uv` or dependencies) are ordered sequentially, with clear warnings on mutually exclusive options. - * **Post-Training Tutorial Integration:** Validates that architectural maps for downstream post-training techniques (GRPO, Distillation, PPO) are integrated without divergence. - ---- - -## 3. Dependency Matrix & Integration Topology - -To guide cross-product testing, the integration pipeline utilizes the following product dependency matrix: - -| Product | Depends On MaxText | Depends On Tunix | Depends On vLLM | Depends On tpu-inference | Depends On xpk | Depends On Pathways | Depends On JAX | -| :--- | :---: | :---: | :---: | :---: | :---: | :---: | :---: | -| **MaxText** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | -| **Tunix** | No | Yes | No | No | No | No | No | -| **vLLM** | No | No | Yes | Yes | Yes | Yes | Yes | -| **tpu-inference** | No | No | Yes | Yes | Yes | Yes | Yes | -| **xpk** | No | No | No | No | Yes | No | Yes | -| **Pathways** | No | No | No | No | No | Yes | Yes | -| **JAX** | No | No | No | No | No | No | Yes | - ---- - -## 4. Phased Implementation Plan - -We propose a four-phase integration roadmap to transition from localized validations to an ecosystem-wide automated CI pipeline: - -### Phase 1: Local AOT Validation & Tutorial Auditing (Weeks 1–2) -* **Action Items:** - * Implement structural mesh dimension checks in `src/maxtext/configs/pyconfig.py` to throw validation errors prior to JAX initialization. - * Integrate standard markdown validators in the documentation build to parse environment variable syntax, check for instruct/non-instruct config consistency, and verify chronological installation steps (e.g., ensuring tools like `uv` are not referenced prior to their setup step). - * Implement path checking in script runners to prevent execution failures when users trigger scripts from unexpected current working directories. - * Expose a command-line dry-run mode: `python3 src/maxtext/train.py base.yml dry_run=True` to compile the logical HLO graph on virtual devices and calculate expected memory. -* **Success Metric:** 100% of incompatible mesh, tutorial path, and prerequisite layout errors fail under 1 second on CPU, before container building or hardware compilation begins. - -### Phase 2: Checkpoint Resharding & Environment Stacking CI (Weeks 3–4) -* **Action Items:** - * Create automated pytest configurations that utilize virtual device grids to simulate resharding layouts. - * Implement dependency isolation checks within setup scripts to verify that installing new packages overrides legacy libraries (`tunix`) without package collision. - * Add pre-execution Docker credentials and service account capability audits to prevent late-stage image push blockages. - * Add standard pre-execution IAM checks in upload scripts to validate registry permission sets. -* **Success Metric:** Checkpoint restoration failures, container push auth blocks, and installation conflicts are caught in sandboxed CPU/virtual meshes prior to hardware allocation. - -### Phase 3: Modular Containers & Verbose Diagnostics (Weeks 5–6) -* **Action Items:** - * Refactor Dockerfile definitions to expose lightweight pathways for SFT and inference setups, eliminating unused package overhead. - * Incorporate standard logging layers, explicit exit codes, and verbose flags into the Docker upload runner utilities, providing time estimates for compilation/build milestones. - * Integrate data-loader step-resume validations and multi-node GRPO connection drop resilience tests into regular integration testing. - * Validate notebook setups against standard TPU VM environments to guarantee that linked templates match hardware profiles. -* **Success Metric:** SFT container compilation times are reduced by at least 2 minutes, and execution failures emit structured logs, progress metrics, and connection retry assertions. - -### Phase 4: Integrated Ecosystem CI & Compatibility Matrix (Weeks 7–8) -* **Action Items:** - * Deploy pre-release hooks that trigger automated testing when dependencies (JAX, vLLM, Tunix) publish Release Candidates. - * Configure a live, auto-updated compatibility dashboard showing verified versions: `MaxText vX.X is compatible with vLLM vY.Y, JAX vZ.Z, and Tunix vA.A`. -* **Success Metric:** 100% of breaking changes introduced by upstream dependencies are caught and notified before the dependency is released to the public. - ---- - -## 5. User Onboarding & First-Run Corner Cases - -When a developer first clones the MaxText repository, they are highly susceptible to environment-related failure modes. The test plan validates the following onboarding corner cases: - -### A. GCS Bucket Write & Permission Mismatches -* **The Failure:** Training compiles, but crashes hours later when attempting the first checkpoint because the environment's Application Default Credentials (ADC) lack write permissions to `base_output_directory`. -* **Validation:** Attempt to write a small, temporary dummy metadata file to the output directory during early startup to verify read/write permissions before graph compilation. - -### B. Gated Hugging Face Repository Credentials -* **The Failure:** Loading gated models (e.g., Llama, Gemma) fails late during execution because the user lacks an active `HF_TOKEN` or hasn't logged in via the Hugging Face CLI. -* **Validation:** Perform an early Hugging Face API request to assert credential validation and access permissions for the selected model family prior to initializing JAX layers. - -### C. Local Storage & Scratch Space Capacity -* **The Failure:** Local runs exhaust disk space when saving large checkpoints, causing process freezes or corrupting existing state. -* **Validation:** Calculate the expected size of checkpoints statically using parameter and optimizer dimensions. Ensure the local target disk has at least `3x` the target checkpoint size in free storage space. - -### D. Physical vs. Configured Device Mesh Mismatches -* **The Failure:** Executing workloads with a configured mesh size that does not align with the physical VM layout results in driver deadlocks. -* **Validation:** Prior to JAX compilation, verify that `prod(ici_fsdp_parallelism, ici_tensor_parallelism, ici_context_parallelism) == jax.local_device_count()`, throwing an immediate topology mismatch exception if they differ. - -### E. Dynamic Data Padding & Recompilation Loops -* **The Failure:** Shifting sequence lengths in input datasets trigger slow, repeated recompilations on every step. -* **Validation:** Monitor initial batch shapes. If sequence sizes vary across the first few steps, throw a diagnostic warning recommending padding configurations or `packing=True`. - -### F. Legacy Package Stacking Collisions -* **The Failure:** Upgraded tutorials or workflows crash due to legacy library versions (`tunix`) lingering in the environment. -* **Validation:** Add a pre-run script check that verifies system packages against active requirements, checking that older library instances are not shadowing the current package installations. - -### G. Stateful Tooling Persistence -* **The Failure:** CLI configurations or temporary elements established for cluster job execution (`xpk`) are cleared out after standard VM restarts. -* **Validation:** Implement persistence verification for configuration caches and tool assets, providing warning diagnostics if essential tools require re-initialization after a VM restart. - -### H. Pre-Flight Docker Authentication Verification -* **The Failure:** Docker image uploads fail with opaque permissions errors at the very end of a long image build because the docker daemon lacks setup credentials. -* **Validation:** Check for active registry endpoints in local Docker config files or run a dry-run credential check before launching build commands. - -### I. Service Account Permissions and Resource Scopes -* **The Failure:** VM workflows crash when retrieving files from GCS or writing images to the Artifact Registry because the attached VM service account lacks required scopes. -* **Validation:** Interrogate local metadata server endpoints during initial setup to assert that the active VM service account possesses correct IAM roles. - -### J. XPK Cluster and Workload Variable Pre-population -* **The Failure:** Submitting cluster tasks through XPK fails because expected variables (TPU type, workload ID, region parameters) are unset, forcing manual cloud console searches. -* **Validation:** Implement an onboarding dry-run script that probes local env variables and suggests export command overrides. - -### K. Notebook Hardware Profiler Sanity Check -* **The Failure:** Users running Colab notebooks hit driver or allocation errors due to the host executing on standard CPU/GPU instances instead of target TPU versions. -* **Validation:** Execute early, non-blocking cell checks to inspect `jax.devices()` and verify that physical accelerator shapes match expected TPU configs (e.g., v5p). - -### L. Distributed Connection Handshakes and Retries -* **The Failure:** Distributed workflows (such as GRPO) encounter transient network dropouts, immediately crashing hours of multi-host execution. -* **Validation:** Configure connection handshake retry budgets and keep-alive settings inside JAX/distributed runtime wrappers to gracefully survive transient drops. - -### M. Standard SSH & VM Package Manager Availability -* **The Failure:** Attempting to run automated setups or SSH overrides fails on stripped custom VM images because package managers (like `apt`) or standard connection paths are absent. -* **Validation:** Pre-flight shell capability checks to confirm minimal command sets are available, providing explicit, manual alternate recipes if executing inside constrained environments. - ---- - -> [!IMPORTANT] -> Integrating AOT validators and cross-dependency hooks directly into the local development and release workflows helps prevent the majority of runtime startup and update-related failures. - -> [!TIP] -> Using JAX's virtual device configurations, these suites can be evaluated within standard CPU-only environments, ensuring high execution velocity and low execution cost. diff --git a/resolve_conflict.py b/resolve_conflict.py deleted file mode 100644 index dd2f53c632..0000000000 --- a/resolve_conflict.py +++ /dev/null @@ -1,23 +0,0 @@ -import sys - -def resolve(): - with open('src/maxtext/layers/quantizations.py', 'r') as f: - lines = f.readlines() - - out = [] - in_conflict = False - for line in lines: - if line.startswith('<<<<<<< HEAD'): - in_conflict = True - continue - if line.startswith('>>>>>>> origin/main'): - in_conflict = False - continue - - if not in_conflict: - out.append(line) - - with open('src/maxtext/layers/quantizations.py', 'w') as f: - f.writelines(out) - -resolve()