diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index b0f407eff3..769a8c1745 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -988,11 +988,11 @@ class HardwareAndMesh(BaseModel): CustomRule.DEFAULT, description="Customized mesh and logical rules for granularity." ) allow_split_physical_axes: bool = Field(False, description="Allow splitting physical axes for device mesh creation.") - enable_nnx: bool = Field(False, description="Whether to use NNX for model definition.") + enable_nnx: bool = Field(True, description="Whether to use NNX for model definition.") optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.") shardy: bool = Field(True, description="Whether to use shardy XLA backend.") - pure_nnx_decoder: bool = Field(False, description="Whether to enable pure NNX decoder.") - pure_nnx: bool = Field(False, description="Whether to enable pure NNX mode.") + pure_nnx_decoder: bool = Field(True, description="Whether to enable pure NNX decoder.") + pure_nnx: bool = Field(True, description="Whether to enable pure NNX mode.") remove_size_one_mesh_axis_from_type: bool = Field( True, description="Whether to remove size one mesh axis from type through jax.config." ) diff --git a/tests/integration/hlo_diff_test.py b/tests/integration/hlo_diff_test.py index cc18713749..baeb8ebbe7 100644 --- a/tests/integration/hlo_diff_test.py +++ b/tests/integration/hlo_diff_test.py @@ -93,11 +93,11 @@ def get_filtered_lines(file_path): return False @pytest.mark.parametrize( - "test_id, config_file, overrides", + "test_id, model_name, overrides", [ ( "deepseek3", - "src/maxtext/configs/models/deepseek3-test.yml", + "deepseek3-test", { "compile_topology": "v6e-4", "base_num_decoder_layers": 4, @@ -107,7 +107,7 @@ def get_filtered_lines(file_path): ), ( "llama3_8b", - "src/maxtext/configs/models/llama3-8b.yml", + "llama3-8b", { "compile_topology": "v6e-4", "base_num_decoder_layers": 4, @@ -117,7 +117,7 @@ def get_filtered_lines(file_path): ), ( "qwen3_1.7b", - "src/maxtext/configs/models/qwen3-1.7b.yml", + "qwen3-1.7b", { "compile_topology": "v6e-4", "base_num_decoder_layers": 4, @@ -127,7 +127,7 @@ def get_filtered_lines(file_path): ), ], ) - def test_hlo_diff(self, test_id, config_file, overrides): + def test_hlo_diff(self, test_id, model_name, overrides): """Test HLO diff for parameterized configurations.""" local_landing_dir = os.path.join(os.path.dirname(__file__), f"hlo_diff_dump_{test_id}") @@ -138,12 +138,16 @@ def test_hlo_diff(self, test_id, config_file, overrides): try: base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) - config_path = os.path.join(base_dir, config_file) + # Load via base.yml + model_name, the normal training path: pyconfig resolves the model + # yml from model_name and merges it onto base.yml. Loading the model yml directly as the + # top-level config skips base.yml, leaving logical_axis_rules empty. + base_config_path = os.path.join(base_dir, "src/maxtext/configs/base.yml") # Arguments for train_compile test_args = [ None, - config_path, + base_config_path, + f"model_name={model_name}", "dataset_type=synthetic", "override_model_config=true", "compile_topology_num_slices=1", diff --git a/tests/utils/reference_hlo_deepseek3.txt b/tests/utils/reference_hlo_deepseek3.txt index eb316f1963..a85276fcd9 100644 --- a/tests/utils/reference_hlo_deepseek3.txt +++ b/tests/utils/reference_hlo_deepseek3.txt @@ -1,4 +1,4 @@ -HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias), {42}: (42, {}, may-alias), {43}: (43, {}, may-alias), {44}: (44, {}, may-alias), {45}: (45, {}, may-alias), {46}: (46, {}, may-alias), {47}: (47, {}, may-alias), {48}: (48, {}, may-alias), {49}: (49, {}, may-alias), {50}: (50, {}, may-alias), {51}: (51, {}, may-alias), {52}: (52, {}, may-alias), {53}: (53, {}, may-alias), {54}: (54, {}, may-alias), {55}: (55, {}, may-alias), {56}: (56, {}, may-alias), {57}: (57, {}, may-alias), {58}: (58, {}, may-alias), {59}: (59, {}, may-alias), {60}: (60, {}, may-alias), {61}: (61, {}, may-alias), {62}: (62, {}, may-alias), {63}: (63, {}, may-alias), {64}: (64, {}, may-alias), {65}: (65, {}, may-alias), {66}: (66, {}, may-alias), {67}: (67, {}, may-alias), {68}: (68, {}, may-alias), {69}: (69, {}, may-alias), {70}: (70, {}, may-alias), {71}: (71, {}, may-alias), {72}: (72, {}, may-alias), {73}: (73, {}, may-alias), {74}: (74, {}, may-alias), {75}: (75, {}, may-alias), {76}: (76, {}, may-alias), {77}: (77, {}, may-alias), {78}: (78, {}, may-alias), {79}: (79, {}, may-alias), {80}: (80, {}, may-alias), {81}: (81, {}, may-alias), {82}: (82, {}, may-alias), {83}: (83, {}, may-alias), {84}: (84, {}, may-alias), {85}: (85, {}, may-alias), {86}: (86, {}, may-alias), {87}: (87, {}, may-alias), {88}: (88, {}, may-alias), {89}: (89, {}, may-alias), {90}: (90, {}, may-alias), {91}: (91, {}, may-alias), {92}: (92, {}, may-alias), {93}: (93, {}, may-alias), {94}: (94, {}, may-alias), {95}: (95, {}, may-alias), {96}: (96, {}, may-alias), {97}: (97, {}, may-alias), {98}: (98, {}, may-alias) }, entry_computation_layout={(s32[]{:T(128)}, f32[512]{0:T(512)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, /*index=5*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=10*/f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, /*index=15*/f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,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[256,1,512,512]{3,2,1,0:T(8,128)}, /*index=20*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=25*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, /*index=30*/f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[512]{0:T(512)}, /*index=35*/f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=40*/f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, /*index=45*/f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,128)}, /*index=50*/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[256,1,512,512]{3,2,1,0:T(8,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, /*index=55*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, /*index=60*/f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, /*index=65*/f32[129280,512]{1,0:T(8,128)}, f32[512]{0:T(512)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, /*index=70*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=75*/f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, /*index=80*/f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,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[256,1,512,512]{3,2,1,0:T(8,128)}, /*index=85*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=90*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, /*index=95*/f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, s32[]{:T(128)}, s32[4,128]{1,0:T(4,128)}, /*index=100*/s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)})->(s32[]{:T(128)}, f32[512]{0:T(512)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, /*index=5*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=10*/f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, /*index=15*/f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,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[256,1,512,512]{3,2,1,0:T(8,128)}, /*index=20*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=25*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, /*index=30*/f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[512]{0:T(512)}, /*index=35*/f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=40*/f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, /*index=45*/f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,128)}, /*index=50*/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[256,1,512,512]{3,2,1,0:T(8,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, /*index=55*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, /*index=60*/f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, /*index=65*/f32[129280,512]{1,0:T(8,128)}, f32[512]{0:T(512)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, /*index=70*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=75*/f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, /*index=80*/f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,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[256,1,512,512]{3,2,1,0:T(8,128)}, /*index=85*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=90*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, /*index=95*/f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[]{:T(128)}, /*index=100*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=105*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 +HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias), {42}: (42, {}, may-alias), {43}: (43, {}, may-alias), {44}: (44, {}, may-alias), {45}: (45, {}, may-alias), {46}: (46, {}, may-alias), {47}: (47, {}, may-alias), {48}: (48, {}, may-alias), {49}: (49, {}, may-alias), {50}: (50, {}, may-alias), {51}: (51, {}, may-alias), {52}: (52, {}, may-alias), {53}: (53, {}, may-alias), {54}: (54, {}, may-alias), {55}: (55, {}, may-alias), {56}: (56, {}, may-alias), {57}: (57, {}, may-alias), {58}: (58, {}, may-alias), {59}: (59, {}, may-alias), {60}: (60, {}, may-alias), {61}: (61, {}, may-alias), {62}: (62, {}, may-alias), {63}: (63, {}, may-alias), {64}: (64, {}, may-alias), {65}: (65, {}, may-alias), {66}: (66, {}, may-alias), {67}: (67, {}, may-alias), {68}: (68, {}, may-alias), {69}: (69, {}, may-alias), {70}: (70, {}, may-alias), {71}: (71, {}, may-alias), {72}: (72, {}, may-alias), {73}: (73, {}, may-alias), {74}: (74, {}, may-alias), {75}: (75, {}, may-alias), {76}: (76, {}, may-alias), {77}: (77, {}, may-alias), {78}: (78, {}, may-alias), {79}: (79, {}, may-alias), {80}: (80, {}, may-alias), {81}: (81, {}, may-alias), {82}: (82, {}, may-alias), {83}: (83, {}, may-alias), {84}: (84, {}, may-alias), {85}: (85, {}, may-alias), {86}: (86, {}, may-alias), {87}: (87, {}, may-alias), {88}: (88, {}, may-alias), {89}: (89, {}, may-alias), {90}: (90, {}, may-alias), {91}: (91, {}, may-alias), {92}: (92, {}, may-alias), {93}: (93, {}, may-alias), {94}: (94, {}, may-alias), {95}: (95, {}, may-alias), {96}: (96, {}, may-alias), {97}: (97, {}, may-alias), {98}: (98, {}, may-alias), {99}: (99, {}, may-alias), {100}: (100, {}, may-alias), {101}: (101, {}, may-alias), {102}: (102, {}, may-alias), {103}: (103, {}, may-alias), {104}: (104, {}, may-alias), {105}: (105, {}, may-alias), {106}: (106, {}, may-alias), {107}: (107, {}, may-alias), {108}: (108, {}, may-alias), {109}: (109, {}, may-alias), {110}: (110, {}, may-alias), {111}: (111, {}, may-alias), {112}: (112, {}, may-alias), {113}: (113, {}, may-alias), {114}: (114, {}, may-alias), {115}: (115, {}, may-alias), {116}: (116, {}, may-alias), {117}: (117, {}, may-alias), {118}: (118, {}, may-alias), {119}: (119, {}, may-alias), {120}: (120, {}, may-alias), {121}: (121, {}, may-alias), {122}: (122, {}, may-alias), {123}: (123, {}, may-alias), {124}: (124, {}, may-alias), {125}: (125, {}, may-alias), {126}: (126, {}, may-alias), {127}: (127, {}, may-alias), {128}: (128, {}, may-alias), {129}: (129, {}, may-alias), {130}: (130, {}, may-alias), {131}: (131, {}, may-alias), {132}: (132, {}, may-alias), {133}: (133, {}, may-alias), {134}: (134, {}, may-alias), {135}: (135, {}, may-alias), {136}: (136, {}, may-alias), {137}: (137, {}, may-alias), {138}: (138, {}, may-alias), {139}: (139, {}, may-alias), {140}: (140, {}, may-alias), {141}: (141, {}, may-alias), {142}: (142, {}, may-alias), {143}: (143, {}, may-alias), {144}: (144, {}, may-alias), {145}: (145, {}, may-alias), {146}: (146, {}, may-alias) }, entry_computation_layout={(f32[512]{0:T(512)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=5*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=10*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, /*index=15*/f32[18432,3,128]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=20*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=25*/f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, /*index=30*/f32[384,3,128,192]{2,3,1,0:T(8,128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=35*/u32[]{:T(128)}, u32[4]{0:T(128)}, f32[128,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, /*index=40*/u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, /*index=45*/u32[1,4]{1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, u32[1]{0:T(128)}, /*index=50*/u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, /*index=55*/f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, /*index=60*/u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=65*/f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[128,1,576]{0,2,1:T(8,128)}, /*index=70*/f32[128,1,128,256]{3,2,1,0:T(8,128)}, f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=75*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, f32[129280,128]{1,0:T(8,128)}, /*index=80*/s32[]{:T(128)}, f32[512]{0:T(512)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, /*index=85*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=90*/f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, /*index=95*/f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, /*index=100*/f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=105*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[128,1,576]{0,2,1:T(8,128)}, f32[128,1,128,256]{3,2,1,0:T(8,128)}, /*index=110*/f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, f32[512]{0:T(512)}, f32[128,3,18432]{2,0,1:T(8,128)}, /*index=115*/f32[128,3,18432]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=120*/f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, /*index=125*/f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, /*index=130*/f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, /*index=135*/f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, /*index=140*/f32[128,1,576]{0,2,1:T(8,128)}, f32[128,1,128,256]{3,2,1,0:T(8,128)}, f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, /*index=145*/s32[]{:T(128)}, u32[]{:T(128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, /*index=150*/s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)})->(f32[512]{0:T(512)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=5*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=10*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, /*index=15*/f32[18432,3,128]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=20*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=25*/f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, /*index=30*/f32[384,3,128,192]{2,3,1,0:T(8,128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=35*/u32[]{:T(128)}, u32[4]{0:T(128)}, f32[128,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, /*index=40*/u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, /*index=45*/u32[1,4]{1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, u32[1]{0:T(128)}, /*index=50*/u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, /*index=55*/f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, /*index=60*/u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=65*/f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[128,1,576]{0,2,1:T(8,128)}, /*index=70*/f32[128,1,128,256]{3,2,1,0:T(8,128)}, f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=75*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, f32[129280,128]{1,0:T(8,128)}, /*index=80*/s32[]{:T(128)}, f32[512]{0:T(512)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, /*index=85*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=90*/f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, /*index=95*/f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, /*index=100*/f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=105*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[128,1,576]{0,2,1:T(8,128)}, f32[128,1,128,256]{3,2,1,0:T(8,128)}, /*index=110*/f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, f32[512]{0:T(512)}, f32[128,3,18432]{2,0,1:T(8,128)}, /*index=115*/f32[128,3,18432]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=120*/f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, /*index=125*/f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, /*index=130*/f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, /*index=135*/f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, /*index=140*/f32[128,1,576]{0,2,1:T(8,128)}, f32[128,1,128,256]{3,2,1,0:T(8,128)}, f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, /*index=145*/s32[]{:T(128)}, u32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=150*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=155*/f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 FileNames @@ -9,31 +9,37 @@ FileLocations StackFrames -%region_46.56 (top_k.25: bf16[], top_k.26: bf16[], top_k.27: s32[], top_k.28: s32[]) -> pred[] { - %constant.1427 = s32[]{:T(128)} constant(0) - %constant.1428 = 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.1427), 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.1428, %bitcast-convert.39), metadata={op_name="xor.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%region_0.1.clone (reduce_sum.820: s32[], reduce_sum.821: s32[]) -> s32[] { + %reduce_sum.820 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.821 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.826 = s32[]{:T(128)} add(%reduce_sum.820, %reduce_sum.821), 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"]}]}} +} + +%region_46.57 (top_k.23: bf16[], top_k.24: bf16[], top_k.25: s32[], top_k.26: s32[]) -> pred[] { + %constant.1441 = s32[]{:T(128)} constant(0) + %constant.1442 = s32[]{:T(128)} constant(2147483647) + %top_k.23 = bf16[]{:T(256)} parameter(0), metadata={op_name="top_k"} + %top_k.24 = bf16[]{:T(256)} parameter(1), metadata={op_name="top_k"} + %top_k.25 = s32[]{:T(128)} parameter(2), metadata={op_name="top_k"} + %top_k.26 = s32[]{:T(128)} parameter(3), metadata={op_name="top_k"} + %convert.392 = f32[]{:T(128)S(6)} convert(%top_k.23), 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.392), 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.1441), 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.1442, %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.1427), 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.1428, %bitcast-convert.40), metadata={op_name="xor.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %convert.393 = f32[]{:T(128)S(6)} convert(%top_k.24), 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.393), 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.1441), 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.1442, %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":[]}} %compare.148 = pred[]{:T(512)S(6)} compare(%compare.146, %compare.147), direction=EQ, metadata={op_name="compare.118"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.149 = pred[]{:T(512)S(6)} compare(%top_k.27, %top_k.28), direction=LT, metadata={op_name="compare.119"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.149 = pred[]{:T(512)S(6)} compare(%top_k.25, %top_k.26), direction=LT, metadata={op_name="compare.119"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} ROOT %select.129 = pred[]{:T(512)} select(%compare.148, %compare.149, %compare.146), metadata={op_name="select.113"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_47.57 (sort.64: s32[], sort.65: s32[], sort.66: s32[], sort.67: s32[]) -> pred[] { +%region_47.58 (sort.64: s32[], sort.65: s32[], sort.66: s32[], sort.67: s32[]) -> pred[] { %sort.64 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(argsort)/sort"} %sort.65 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(argsort)/sort"} %sort.66 = s32[]{:T(128)} parameter(2), metadata={op_name="jit(argsort)/sort"} @@ -45,7 +51,7 @@ StackFrames ROOT %select.130 = pred[]{:T(512)} select(%compare.150, %compare.151, %lt_to.32), metadata={op_name="select.114"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_48.58 (sort.68: s32[], sort.69: s32[], sort.70: s32[], sort.71: s32[]) -> pred[] { +%region_48.59 (sort.68: s32[], sort.69: s32[], sort.70: s32[], sort.71: s32[]) -> pred[] { %sort.68 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(argsort)/sort"} %sort.69 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(argsort)/sort"} %sort.70 = s32[]{:T(128)} parameter(2), metadata={op_name="jit(argsort)/sort"} @@ -57,7 +63,7 @@ StackFrames ROOT %select.131 = pred[]{:T(512)} select(%compare.152, %compare.153, %lt_to.34), metadata={op_name="select.115"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_67.80 (sort.78: s32[], sort.79: s32[], sort.80: s32[], sort.81: s32[]) -> pred[] { +%region_67.81 (sort.78: s32[], sort.79: s32[], sort.80: s32[], sort.81: s32[]) -> pred[] { %sort.78 = s32[]{:T(128)} parameter(0), metadata={op_name="sort_activations/jit(argsort)/sort"} %sort.79 = s32[]{:T(128)} parameter(1), metadata={op_name="sort_activations/jit(argsort)/sort"} %sort.80 = s32[]{:T(128)} parameter(2), metadata={op_name="sort_activations/jit(argsort)/sort"} @@ -69,309 +75,315 @@ StackFrames ROOT %select.134 = pred[]{:T(512)} select(%compare.156, %compare.157, %lt_to.37), metadata={op_name="select.116"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_119.141 (reduce_sum.157: bf16[], reduce_sum.158: bf16[]) -> bf16[] { - %reduce_sum.157 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/reduce_sum"} - %reduce_sum.158 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/reduce_sum"} - ROOT %reduce_sum.159 = bf16[]{:T(256)} add(%reduce_sum.157, %reduce_sum.158), metadata={op_name="checkpoint/moe_layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_107.126 (psum.6: bf16[], psum.9: bf16[]) -> bf16[] { +%region_107.127 (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.1445 = 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.1446 = 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.1450 = 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_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.1447 = 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[] { +%region_62.74 (reduce-window.111: s32[], reduce-window.112: s32[]) -> s32[] { %reduce-window.111 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.35"} %reduce-window.112 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.35"} - ROOT %reduce_window_sum.108 = s32[]{:T(128)} add(%reduce-window.111, %reduce-window.112), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.108 = s32[]{:T(128)} add(%reduce-window.111, %reduce-window.112), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_64.75 (reduce-window.113: s32[], reduce-window.114: s32[]) -> s32[] { +%region_64.76 (reduce-window.113: s32[], reduce-window.114: s32[]) -> s32[] { %reduce-window.113 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.36"} %reduce-window.114 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.36"} - ROOT %reduce_window_sum.109 = s32[]{:T(128)} add(%reduce-window.113, %reduce-window.114), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.109 = s32[]{:T(128)} add(%reduce-window.113, %reduce-window.114), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_65.76 (reduce-window.115: s32[], reduce-window.116: s32[]) -> s32[] { +%region_65.77 (reduce-window.115: s32[], reduce-window.116: s32[]) -> s32[] { %reduce-window.115 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.63"} %reduce-window.116 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.63"} - ROOT %reduce_window_sum.110 = s32[]{:T(128)} add(%reduce-window.115, %reduce-window.116), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.110 = s32[]{:T(128)} add(%reduce-window.115, %reduce-window.116), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_68.81.clone (reduce-window.396: s32[], reduce-window.397: s32[]) -> s32[] { +%region_68.82.clone (reduce-window.396: s32[], reduce-window.397: s32[]) -> s32[] { %reduce-window.396 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.38"} %reduce-window.397 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.38"} - ROOT %reduce_window_sum.317 = s32[]{:T(128)} add(%reduce-window.396, %reduce-window.397), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.317 = s32[]{:T(128)} add(%reduce-window.396, %reduce-window.397), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_69.82.clone (reduce-window.400: s32[], reduce-window.401: s32[]) -> s32[] { +%region_69.83.clone (reduce-window.400: s32[], reduce-window.401: s32[]) -> s32[] { %reduce-window.400 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.64"} %reduce-window.401 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.64"} - ROOT %reduce_window_sum.319 = s32[]{:T(128)} add(%reduce-window.400, %reduce-window.401), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.319 = s32[]{:T(128)} add(%reduce-window.400, %reduce-window.401), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_71.84.clone (reduce-window.404: s32[], reduce-window.405: s32[]) -> s32[] { +%region_71.85.clone (reduce-window.404: s32[], reduce-window.405: s32[]) -> s32[] { %reduce-window.404 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.40"} %reduce-window.405 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.40"} - ROOT %reduce_window_sum.321 = s32[]{:T(128)} add(%reduce-window.404, %reduce-window.405), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.321 = s32[]{:T(128)} add(%reduce-window.404, %reduce-window.405), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_72.85.clone (reduce-window.408: s32[], reduce-window.409: s32[]) -> s32[] { +%region_72.86.clone (reduce-window.408: s32[], reduce-window.409: s32[]) -> s32[] { %reduce-window.408 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.65"} %reduce-window.409 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.65"} - ROOT %reduce_window_sum.323 = s32[]{:T(128)} add(%reduce-window.408, %reduce-window.409), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.323 = s32[]{:T(128)} add(%reduce-window.408, %reduce-window.409), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_74.87.clone (reduce-window.412: s32[], reduce-window.413: s32[]) -> s32[] { +%region_74.88.clone (reduce-window.412: s32[], reduce-window.413: s32[]) -> s32[] { %reduce-window.412 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.42"} %reduce-window.413 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.42"} - ROOT %reduce_window_sum.325 = s32[]{:T(128)} add(%reduce-window.412, %reduce-window.413), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.325 = s32[]{:T(128)} add(%reduce-window.412, %reduce-window.413), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_75.88.clone (reduce-window.416: s32[], reduce-window.417: s32[]) -> s32[] { +%region_75.89.clone (reduce-window.416: s32[], reduce-window.417: s32[]) -> s32[] { %reduce-window.416 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.66"} %reduce-window.417 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.66"} - ROOT %reduce_window_sum.327 = s32[]{:T(128)} add(%reduce-window.416, %reduce-window.417), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.327 = s32[]{:T(128)} add(%reduce-window.416, %reduce-window.417), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_80.96.clone (reduce-window.420: s32[], reduce-window.421: s32[]) -> s32[] { +%region_80.97.clone (reduce-window.420: s32[], reduce-window.421: s32[]) -> s32[] { %reduce-window.420 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.44"} %reduce-window.421 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.44"} - ROOT %reduce_window_sum.329 = s32[]{:T(128)} add(%reduce-window.420, %reduce-window.421), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.329 = s32[]{:T(128)} add(%reduce-window.420, %reduce-window.421), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_82.98.clone (reduce-window.424: s32[], reduce-window.425: s32[]) -> s32[] { +%region_82.99.clone (reduce-window.424: s32[], reduce-window.425: s32[]) -> s32[] { %reduce-window.424 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.45"} %reduce-window.425 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.45"} - ROOT %reduce_window_sum.331 = s32[]{:T(128)} add(%reduce-window.424, %reduce-window.425), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.331 = s32[]{:T(128)} add(%reduce-window.424, %reduce-window.425), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_83.99.clone (reduce-window.428: s32[], reduce-window.429: s32[]) -> s32[] { +%region_83.100.clone (reduce-window.428: s32[], reduce-window.429: s32[]) -> s32[] { %reduce-window.428 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.67"} %reduce-window.429 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.67"} - ROOT %reduce_window_sum.333 = s32[]{:T(128)} add(%reduce-window.428, %reduce-window.429), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.333 = s32[]{:T(128)} add(%reduce-window.428, %reduce-window.429), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_94.112 (reduce-window.174: s32[], reduce-window.175: s32[]) -> s32[] { +%region_92.111 (reduce-window.174: s32[], reduce-window.175: s32[]) -> s32[] { %reduce-window.174 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.49"} %reduce-window.175 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.49"} - ROOT %reduce_window_sum.138 = s32[]{:T(128)} add(%reduce-window.174, %reduce-window.175), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.138 = s32[]{:T(128)} add(%reduce-window.174, %reduce-window.175), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_95.113 (reduce-window.179: s32[], reduce-window.180: s32[]) -> s32[] { +%region_93.112 (reduce-window.179: s32[], reduce-window.180: s32[]) -> s32[] { %reduce-window.179 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.69"} %reduce-window.180 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.69"} - ROOT %reduce_window_sum.139 = s32[]{:T(128)} add(%reduce-window.179, %reduce-window.180), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.139 = s32[]{:T(128)} add(%reduce-window.179, %reduce-window.180), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_97.115 (reduce-window.184: s32[], reduce-window.185: s32[]) -> s32[] { +%region_95.114 (reduce-window.184: s32[], reduce-window.185: s32[]) -> s32[] { %reduce-window.184 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.51"} %reduce-window.185 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.51"} - ROOT %reduce_window_sum.140 = s32[]{:T(128)} add(%reduce-window.184, %reduce-window.185), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.140 = s32[]{:T(128)} add(%reduce-window.184, %reduce-window.185), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_98.116 (reduce-window.189: s32[], reduce-window.190: s32[]) -> s32[] { +%region_96.115 (reduce-window.189: s32[], reduce-window.190: s32[]) -> s32[] { %reduce-window.189 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.70"} %reduce-window.190 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.70"} - ROOT %reduce_window_sum.141 = s32[]{:T(128)} add(%reduce-window.189, %reduce-window.190), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.141 = s32[]{:T(128)} add(%reduce-window.189, %reduce-window.190), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_103.121 (reduce-window.194: s32[], reduce-window.195: s32[]) -> s32[] { +%region_101.120 (reduce-window.194: s32[], reduce-window.195: s32[]) -> s32[] { %reduce-window.194 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.53"} %reduce-window.195 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.53"} - ROOT %reduce_window_sum.142 = s32[]{:T(128)} add(%reduce-window.194, %reduce-window.195), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.142 = s32[]{:T(128)} add(%reduce-window.194, %reduce-window.195), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_105.123 (reduce-window.199: s32[], reduce-window.200: s32[]) -> s32[] { +%region_103.122 (reduce-window.199: s32[], reduce-window.200: s32[]) -> s32[] { %reduce-window.199 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.54"} %reduce-window.200 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.54"} - ROOT %reduce_window_sum.143 = s32[]{:T(128)} add(%reduce-window.199, %reduce-window.200), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.143 = s32[]{:T(128)} add(%reduce-window.199, %reduce-window.200), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%region_106.124 (reduce-window.204: s32[], reduce-window.205: s32[]) -> s32[] { +%region_104.123 (reduce-window.204: s32[], reduce-window.205: s32[]) -> s32[] { %reduce-window.204 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.71"} %reduce-window.205 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.71"} - ROOT %reduce_window_sum.144 = s32[]{:T(128)} add(%reduce-window.204, %reduce-window.205), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.5 (param_0.17: bf16[129280,512], param_1.108: s32[1024]) -> bf16[512,512] { - %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.3306 = 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.3306), 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.3305 = 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_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.3314 = 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.3314), 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.3313 = 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] { + ROOT %reduce_window_sum.144 = s32[]{:T(128)} add(%reduce-window.204, %reduce-window.205), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} +} + +%fused_computation.6 (param_0.20: bf16[129280,128], param_1.121: s32[1024]) -> bf16[512,128] { + %param_0.20 = bf16[129280,128]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.121 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.18 = s32[1024]{0:T(1024)} custom-call(%param_1.121), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %slice.916 = s32[512]{0:T(512)} slice(%custom-call.18), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %reshape.4235 = s32[4,128]{1,0:T(4,128)} reshape(%slice.916), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.889 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.4235), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %gather.200 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)} gather(%param_0.20, %transpose.889), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,128}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.888 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)} transpose(%gather.200), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %reshape.4234 = bf16[512,128]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.888), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} +} + +%fused_computation.7 (param_0.23: f32[163840,32], param_1.123: s32[1024]) -> f32[128,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.3322 = 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.3322), 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.3321 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.858), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} + %param_1.123 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.20 = s32[1024]{0:T(1024)} custom-call(%param_1.123), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %slice.918 = s32[128]{0:T(128)} slice(%custom-call.20), slice={[0:128]}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %reshape.4243 = s32[128]{0:T(128)} reshape(%slice.918), metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %transpose.895 = s32[128]{0:T(128)} transpose(%reshape.4243), dimensions={0}, metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %gather.202 = f32[128,32]{1,0:T(8,128)} gather(%param_0.23, %transpose.895), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,32}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %transpose.894 = f32[128,32]{1,0:T(8,128)} transpose(%gather.202), dimensions={0,1}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + ROOT %reshape.4242 = f32[128,32]{1,0:T(8,128)S(1)} reshape(%transpose.894), metadata={op_name="jit(train_step)/gather" stack_frame_id=0} } -%fused_computation.8 (param_0.26: f32[163840,32], param_1.120: s32[1024]) -> f32[512,32] { +%fused_computation.8 (param_0.26: f32[163840,32], param_1.125: s32[1024]) -> f32[128,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.3330 = 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.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.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.3329 = 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} + %param_1.125 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.22 = s32[1024]{0:T(1024)} custom-call(%param_1.125), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %slice.920 = s32[128]{0:T(128)} slice(%custom-call.22), slice={[0:128]}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %reshape.4251 = s32[128]{0:T(128)} reshape(%slice.920), metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %transpose.901 = s32[128]{0:T(128)} transpose(%reshape.4251), dimensions={0}, metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %gather.204 = f32[128,32]{1,0:T(8,128)} gather(%param_0.26, %transpose.901), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,32}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %transpose.900 = f32[128,32]{1,0:T(8,128)} transpose(%gather.204), dimensions={0,1}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + ROOT %reshape.4250 = f32[128,32]{1,0:T(8,128)S(1)} reshape(%transpose.900), metadata={op_name="jit(train_step)/gather" stack_frame_id=0} } -%fused_computation.9 (param_0.29: f32[163840,32], param_1.122: s32[1024]) -> f32[512,32] { +%fused_computation.9 (param_0.29: f32[163840,32], param_1.133: s32[1024]) -> f32[128,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.3338 = 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.3338), 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.3337 = 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.3346 = 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.3346), 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.3345 = 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] { - %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.3354 = 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.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.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.3353 = 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.3362 = 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.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.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.3361 = 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.3370 = 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.3370), 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.3369 = 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.3401 = 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.3401), 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.3400 = 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.3424 = 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.3424), 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.3423 = 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.1918 = 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.3478 = 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.3478), 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.3479 = 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.3479), 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} + %param_1.133 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.30 = s32[1024]{0:T(1024)} custom-call(%param_1.133), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %slice.928 = s32[128]{0:T(128)} slice(%custom-call.30), slice={[0:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %reshape.4259 = s32[128]{0:T(128)} reshape(%slice.928), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + %transpose.907 = s32[128]{0:T(128)} transpose(%reshape.4259), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + %gather.206 = f32[128,32]{1,0:T(8,128)} gather(%param_0.29, %transpose.907), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %transpose.906 = f32[128,32]{1,0:T(8,128)} transpose(%gather.206), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + ROOT %reshape.4258 = f32[128,32]{1,0:T(8,128)} reshape(%transpose.906), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} +} + +%fused_computation.10 (param_0.32: f32[163840,32], param_1.135: s32[1024]) -> f32[128,32] { + %param_0.32 = f32[163840,32]{1,0:T(8,128)} parameter(0) + %param_1.135 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.32 = s32[1024]{0:T(1024)} custom-call(%param_1.135), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %slice.930 = s32[128]{0:T(128)} slice(%custom-call.32), slice={[0:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %reshape.4267 = s32[128]{0:T(128)} reshape(%slice.930), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + %transpose.913 = s32[128]{0:T(128)} transpose(%reshape.4267), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + %gather.208 = f32[128,32]{1,0:T(8,128)} gather(%param_0.32, %transpose.913), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %transpose.912 = f32[128,32]{1,0:T(8,128)} transpose(%gather.208), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + ROOT %reshape.4266 = f32[128,32]{1,0:T(8,128)} reshape(%transpose.912), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} +} + +%fused_computation.11 (param_0.35: bf16[1024,512], param_1.139: s32[1024]) -> bf16[1024,512] { + %param_0.35 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.139 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.36 = s32[1024]{0:T(1024)} custom-call(%param_1.139), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.934 = s32[1024]{0:T(1024)} slice(%custom-call.36), slice={[0:1024]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.4275 = s32[1024]{0:T(1024)} reshape(%slice.934), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.919 = s32[1024]{0:T(1024)} transpose(%reshape.4275), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.210 = bf16[1024,512]{1,0:T(8,128)(2,1)} gather(%param_0.35, %transpose.919), 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())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.918 = bf16[1024,512]{1,0:T(8,128)(2,1)} transpose(%gather.210), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.4274 = bf16[1024,512]{1,0:T(8,128)(2,1)} reshape(%transpose.918), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} +} + +%fused_computation.12 (param_0.38: bf16[128,256], param_1.44: s32[1024]) -> bf16[1024] { + %param_0.38 = bf16[128,256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.3895 = s32[]{:T(128)} constant(127), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %broadcast.3457 = s32[128,8,1]{2,1,0:T(8,128)} broadcast(%constant.3895), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %param_1.44 = s32[1024]{0:T(1024)S(1)} parameter(1) + %reshape.4283 = s32[128,8,1]{2,1,0:T(8,128)} reshape(%param_1.44), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %transpose.925 = s32[128,8,1]{2,1,0:T(8,128)} transpose(%reshape.4283), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %and.347 = s32[128,8,1]{2,1,0:T(8,128)} and(%broadcast.3457, %transpose.925), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %constant.3894 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %broadcast.3456 = s32[128,8,1]{2,1,0:T(8,128)} broadcast(%constant.3894), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %shift-right-logical.22 = s32[128,8,1]{2,1,0:T(8,128)} shift-right-logical(%and.347, %broadcast.3456), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %constant.3893 = s32[]{:T(128)} constant(32640), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %broadcast.3455 = s32[128,8,1]{2,1,0:T(8,128)} broadcast(%constant.3893), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %and.346 = s32[128,8,1]{2,1,0:T(8,128)} and(%broadcast.3455, %transpose.925), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %constant.3892 = s32[]{:T(128)} constant(7), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %broadcast.3454 = s32[128,8,1]{2,1,0:T(8,128)} broadcast(%constant.3892), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %shift-right-logical.21 = s32[128,8,1]{2,1,0:T(8,128)} shift-right-logical(%and.346, %broadcast.3454), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %concatenate.405 = s32[128,8,2]{2,1,0:T(8,128)} concatenate(%shift-right-logical.22, %shift-right-logical.21), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %custom-call.3 = s32[128,8,2]{2,1,0:T(8,128)} custom-call(%concatenate.405), custom_call_target="GatherScatterIndicesBitpacked", metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %gather.212 = bf16[128,8]{1,0:T(8,128)(2,1)} gather(%param_0.38, %custom-call.3), offset_dims={}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %transpose.924 = bf16[128,8]{1,0:T(8,128)(2,1)} transpose(%gather.212), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + ROOT %reshape.4282 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} reshape(%transpose.924), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} +} + +%fused_computation.13 (param_0.41: bf16[1024,512], param_1.141: s32[1024]) -> bf16[1024,512] { + %param_0.41 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.141 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.38 = s32[1024]{0:T(1024)} custom-call(%param_1.141), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.936 = s32[1024]{0:T(1024)} slice(%custom-call.38), slice={[0:1024]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.4291 = s32[1024]{0:T(1024)} reshape(%slice.936), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.931 = s32[1024]{0:T(1024)} transpose(%reshape.4291), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.214 = bf16[1024,512]{1,0:T(8,128)(2,1)} gather(%param_0.41, %transpose.931), 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())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.930 = bf16[1024,512]{1,0:T(8,128)(2,1)} transpose(%gather.214), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.4290 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.930), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} +} + +%fused_computation.14 (param_0.44: bf16[1024,512], param_1.143: s32[1024]) -> bf16[1024,512] { + %param_0.44 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.143 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.40 = s32[1024]{0:T(1024)} custom-call(%param_1.143), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.938 = s32[1024]{0:T(1024)} slice(%custom-call.40), slice={[0:1024]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.4299 = s32[1024]{0:T(1024)} reshape(%slice.938), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.937 = s32[1024]{0:T(1024)} transpose(%reshape.4299), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.216 = bf16[1024,512]{1,0:T(8,128)(2,1)} gather(%param_0.44, %transpose.937), 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())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.936 = bf16[1024,512]{1,0:T(8,128)(2,1)} transpose(%gather.216), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.4298 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.936), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/gather" stack_frame_id=0} +} + +%fused_computation.15 (param_0.47: bf16[1024,512], param_1.145: s32[1024]) -> bf16[1024,512] { + %param_0.47 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.145 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.42 = s32[1024]{0:T(1024)} custom-call(%param_1.145), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.940 = s32[1024]{0:T(1024)} slice(%custom-call.42), slice={[0:1024]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.4307 = s32[1024]{0:T(1024)} reshape(%slice.940), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.943 = s32[1024]{0:T(1024)} transpose(%reshape.4307), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.218 = bf16[1024,512]{1,0:T(8,128)(2,1)} gather(%param_0.47, %transpose.943), 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())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.942 = bf16[1024,512]{1,0:T(8,128)(2,1)} transpose(%gather.218), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.4306 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.942), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/sort_activations/gather" stack_frame_id=0} +} + +%fused_computation.17 (param_0.53: s32[256], param_1.137: s32[1024]) -> s32[257] { + %param_0.53 = s32[256]{0:T(256)S(1)} parameter(0) + %param_1.137 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.34 = s32[1024]{0:T(1024)} custom-call(%param_1.137), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %slice.932 = s32[257]{0:T(512)} slice(%custom-call.34), slice={[0:257]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.4332 = s32[257]{0:T(512)} reshape(%slice.932), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.959 = s32[257]{0:T(512)} transpose(%reshape.4332), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.223 = s32[257]{0:T(512)} gather(%param_0.53, %transpose.959), 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())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.958 = s32[257]{0:T(512)} transpose(%gather.223), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.4331 = s32[257]{0:T(512)S(1)} reshape(%transpose.958), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} +} + +%fused_computation.18 (param_0.56: s32[256], param_1.147: s32[1024]) -> s32[257] { + %param_0.56 = s32[256]{0:T(256)S(1)} parameter(0) + %param_1.147 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.44 = s32[1024]{0:T(1024)} custom-call(%param_1.147), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %slice.942 = s32[257]{0:T(512)} slice(%custom-call.44), slice={[0:257]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.4349 = s32[257]{0:T(512)} reshape(%slice.942), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.969 = s32[257]{0:T(512)} transpose(%reshape.4349), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.226 = s32[257]{0:T(512)} gather(%param_0.56, %transpose.969), 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())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.968 = s32[257]{0:T(512)} transpose(%gather.226), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.4348 = s32[257]{0:T(512)S(1)} reshape(%transpose.968), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} +} + +%region_163.189.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())/scatter-add"} + %scatter-add.96 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add"} + ROOT %add.1991 = bf16[]{:T(256)} add(%scatter-add.94, %scatter-add.96), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.23 (param_0.61: bf16[129280,128], param_1.77: s32[512], param_2.38: bf16[512,128]) -> bf16[129280,128] { + %param_0.61 = bf16[129280,128]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.77 = s32[512]{0:T(512)S(1)} parameter(1) + %reshape.4397 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.77), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %transpose.1002 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.4397), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %param_2.38 = bf16[512,128]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %reshape.4398 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)} reshape(%param_2.38), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %transpose.1003 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)} transpose(%reshape.4398), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + ROOT %scatter.73 = bf16[129280,128]{1,0:T(8,128)(2,1)S(1)} scatter(%param_0.61, %transpose.1002, %transpose.1003), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_163.189.clone, metadata={op_name="jit(train_step)/transpose(jvp())/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.1387 = s32[]{:T(128)} constant(0) - %constant.1388 = s32[]{:T(128)} constant(2147483647) + %constant.1399 = s32[]{:T(128)} constant(0) + %constant.1400 = 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.1387), 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.1388, %bitcast-convert.35), metadata={op_name="xor.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %convert.384 = 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.384), 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.1399), 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.1400, %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.1387), 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.1388, %bitcast-convert.36), metadata={op_name="xor.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %convert.385 = 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.385), 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.1399), 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.1400, %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":[]}} @@ -383,67 +395,67 @@ StackFrames %region_15.21.clone.1 (reduce-window.326: s32[], reduce-window.327: s32[]) -> s32[] { %reduce-window.326 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.20"} %reduce-window.327 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.20"} - ROOT %reduce_window_sum.282 = s32[]{:T(128)} add(%reduce-window.326, %reduce-window.327), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.282 = s32[]{:T(128)} add(%reduce-window.326, %reduce-window.327), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } %region_16.22.clone.1 (reduce-window.330: s32[], reduce-window.331: s32[]) -> s32[] { %reduce-window.330 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.56"} %reduce-window.331 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.56"} - ROOT %reduce_window_sum.284 = s32[]{:T(128)} add(%reduce-window.330, %reduce-window.331), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.284 = s32[]{:T(128)} add(%reduce-window.330, %reduce-window.331), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } %region_18.24.clone.1 (reduce-window.334: s32[], reduce-window.335: s32[]) -> s32[] { %reduce-window.334 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.22"} %reduce-window.335 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.22"} - ROOT %reduce_window_sum.286 = s32[]{:T(128)} add(%reduce-window.334, %reduce-window.335), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.286 = s32[]{:T(128)} add(%reduce-window.334, %reduce-window.335), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } %region_19.25.clone.1 (reduce-window.338: s32[], reduce-window.339: s32[]) -> s32[] { %reduce-window.338 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.57"} %reduce-window.339 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.57"} - ROOT %reduce_window_sum.288 = s32[]{:T(128)} add(%reduce-window.338, %reduce-window.339), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.288 = s32[]{:T(128)} add(%reduce-window.338, %reduce-window.339), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } %region_21.27.clone.1 (reduce-window.342: s32[], reduce-window.343: s32[]) -> s32[] { %reduce-window.342 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.24"} %reduce-window.343 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.24"} - ROOT %reduce_window_sum.290 = s32[]{:T(128)} add(%reduce-window.342, %reduce-window.343), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.290 = s32[]{:T(128)} add(%reduce-window.342, %reduce-window.343), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } %region_22.28.clone.1 (reduce-window.346: s32[], reduce-window.347: s32[]) -> s32[] { %reduce-window.346 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.58"} %reduce-window.347 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.58"} - ROOT %reduce_window_sum.292 = s32[]{:T(128)} add(%reduce-window.346, %reduce-window.347), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.292 = s32[]{:T(128)} add(%reduce-window.346, %reduce-window.347), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } -%fused_computation.4.clone (param_0.68: s32[256], param_1.114: s32[1024]) -> s32[263] { - %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.3622 = 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.3622), 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.3621 = 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} +%fused_computation.5.clone (param_0.74: s32[256], param_1.127: s32[1024]) -> s32[257] { + %param_0.74 = s32[256]{0:T(256)S(1)} parameter(0) + %param_1.127 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.24 = s32[1024]{0:T(1024)} custom-call(%param_1.127), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %slice.922 = s32[257]{0:T(512)} slice(%custom-call.24), slice={[0:257]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %reshape.4623 = s32[257]{0:T(512)} reshape(%slice.922), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %transpose.1085 = s32[257]{0:T(512)} transpose(%reshape.4623), dimensions={0}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} + %gather.228 = s32[257]{0:T(512)} gather(%param_0.74, %transpose.1085), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + %transpose.1084 = s32[257]{0:T(512)} transpose(%gather.228), dimensions={0}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} + ROOT %reshape.4622 = s32[257]{0:T(512)S(1)} reshape(%transpose.1084), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/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[] { %reduce-window.350 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.26"} %reduce-window.351 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.26"} - ROOT %reduce_window_sum.294 = s32[]{:T(128)} add(%reduce-window.350, %reduce-window.351), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.294 = s32[]{:T(128)} add(%reduce-window.350, %reduce-window.351), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } %region_29.36.clone.1 (reduce-window.354: s32[], reduce-window.355: s32[]) -> s32[] { %reduce-window.354 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.27"} %reduce-window.355 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.27"} - ROOT %reduce_window_sum.296 = s32[]{:T(128)} add(%reduce-window.354, %reduce-window.355), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.296 = s32[]{:T(128)} add(%reduce-window.354, %reduce-window.355), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } %region_30.37.clone.1 (reduce-window.358: s32[], reduce-window.359: s32[]) -> s32[] { %reduce-window.358 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.59"} %reduce-window.359 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.59"} - ROOT %reduce_window_sum.298 = s32[]{:T(128)} add(%reduce-window.358, %reduce-window.359), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %reduce_window_sum.298 = s32[]{:T(128)} add(%reduce-window.358, %reduce-window.359), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} } %region_13.19 (sort.44: s32[], sort.45: s32[], sort.46: s32[], sort.47: s32[], sort.48: s32[], sort.49: s32[]) -> pred[] { @@ -460,16 +472,16 @@ StackFrames ROOT %select.121 = pred[]{:T(512)} select(%compare.134, %compare.135, %lt_to.27), metadata={op_name="select.109"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.2.clone (param_0.71: bf16[4096,512], param_1.116: s32[4096]) -> bf16[4096,512] { - %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.3645 = 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.3645), 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.3644 = 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} +%fused_computation.2.clone (param_0.77: bf16[1024,512], param_1.129: s32[1024]) -> bf16[1024,512] { + %param_0.77 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.129 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.26 = s32[1024]{0:T(1024)} custom-call(%param_1.129), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.924 = s32[1024]{0:T(1024)} slice(%custom-call.26), slice={[0:1024]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.4639 = s32[1024]{0:T(1024)} reshape(%slice.924), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.1091 = s32[1024]{0:T(1024)} transpose(%reshape.4639), dimensions={0}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.229 = bf16[1024,512]{1,0:T(8,128)(2,1)} gather(%param_0.77, %transpose.1091), 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()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.1090 = bf16[1024,512]{1,0:T(8,128)(2,1)} transpose(%gather.229), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.4638 = bf16[1024,512]{1,0:T(8,128)(2,1)} reshape(%transpose.1090), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/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[] { @@ -486,16 +498,40 @@ StackFrames ROOT %select.126 = pred[]{:T(512)} select(%compare.142, %compare.143, %lt_to.30), metadata={op_name="select.110"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.3.clone (param_0.72: bf16[4096,512], param_1.118: s32[4096]) -> bf16[4096,512] { - %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.3647 = 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.3647), 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.3646 = 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} +%fused_computation.3.clone (param_0.78: bf16[1024,512], param_1.131: s32[1024]) -> bf16[1024,512] { + %param_0.78 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.131 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.28 = s32[1024]{0:T(1024)} custom-call(%param_1.131), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} + %slice.926 = s32[1024]{0:T(1024)} slice(%custom-call.28), slice={[0:1024]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} + %reshape.4641 = s32[1024]{0:T(1024)} reshape(%slice.926), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %transpose.1093 = s32[1024]{0:T(1024)} transpose(%reshape.4641), dimensions={0}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} + %gather.230 = bf16[1024,512]{1,0:T(8,128)(2,1)} gather(%param_0.78, %transpose.1093), 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()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} + %transpose.1092 = bf16[1024,512]{1,0:T(8,128)(2,1)} transpose(%gather.230), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} + ROOT %reshape.4640 = bf16[1024,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1092), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/sort_activations/gather" stack_frame_id=0} +} + +%fused_computation.4.clone (param_0.79: bf16[128,256], param_1.115: s32[1024]) -> bf16[1024] { + %param_0.79 = bf16[128,256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.3938 = s32[]{:T(128)} constant(127), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %broadcast.3552 = s32[128,8,1]{2,1,0:T(8,128)} broadcast(%constant.3938), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %param_1.115 = s32[1024]{0:T(1024)S(1)} parameter(1) + %reshape.4646 = s32[128,8,1]{2,1,0:T(8,128)} reshape(%param_1.115), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %transpose.1095 = s32[128,8,1]{2,1,0:T(8,128)} transpose(%reshape.4646), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %and.360 = s32[128,8,1]{2,1,0:T(8,128)} and(%broadcast.3552, %transpose.1095), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %constant.3937 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %broadcast.3551 = s32[128,8,1]{2,1,0:T(8,128)} broadcast(%constant.3937), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %shift-right-logical.24 = s32[128,8,1]{2,1,0:T(8,128)} shift-right-logical(%and.360, %broadcast.3551), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %constant.3936 = s32[]{:T(128)} constant(32640), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %broadcast.3550 = s32[128,8,1]{2,1,0:T(8,128)} broadcast(%constant.3936), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %and.359 = s32[128,8,1]{2,1,0:T(8,128)} and(%broadcast.3550, %transpose.1095), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %constant.3935 = s32[]{:T(128)} constant(7), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %broadcast.3549 = s32[128,8,1]{2,1,0:T(8,128)} broadcast(%constant.3935), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %shift-right-logical.23 = s32[128,8,1]{2,1,0:T(8,128)} shift-right-logical(%and.359, %broadcast.3549), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %concatenate.415 = s32[128,8,2]{2,1,0:T(8,128)} concatenate(%shift-right-logical.24, %shift-right-logical.23), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %custom-call.4 = s32[128,8,2]{2,1,0:T(8,128)} custom-call(%concatenate.415), custom_call_target="GatherScatterIndicesBitpacked", metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %gather.231 = bf16[128,8]{1,0:T(8,128)(2,1)} gather(%param_0.79, %custom-call.4), offset_dims={}, collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=2, slice_sizes={1,1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + %transpose.1094 = bf16[128,8]{1,0:T(8,128)(2,1)} transpose(%gather.231), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} + ROOT %reshape.4645 = bf16[1024]{0:T(1024)(128)(2,1)S(1)} reshape(%transpose.1094), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(take_along_axis)/gather" stack_frame_id=0} } %compare (name: s32[], name.1: s32[], name.2: bf16[], name.3: bf16[]) -> pred[] { @@ -503,7 +539,87 @@ 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.293 = pred[] compare(%name, %name.1), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%add.39.clone (x.79: bf16[], y.79: bf16[]) -> bf16[] { + %x.79 = bf16[]{:T(256)} parameter(0) + %y.79 = bf16[]{:T(256)} parameter(1) + ROOT %add.1990 = bf16[]{:T(256)} add(%x.79, %y.79), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%all-reduce-scatter.4 (input.4: bf16[512,129280]) -> bf16[128,129280] { + %input.4 = bf16[512,129280]{1,0:T(8,128)(2,1)} parameter(0) + %all-reduce.144 = bf16[512,129280]{1,0:T(8,128)(2,1)} all-reduce(%input.4), channel_id=171, replica_groups={{0,1,2,3}}, use_global_device_ids=true, to_apply=%add.39.clone, frontend_attributes={from-cross-replica-sharding="true"}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"2"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + %partition-id.63 = u32[] partition-id() + %constant.4159 = u32[] constant(128) + %multiply.1330 = u32[]{:T(128)} multiply(%partition-id.63, %constant.4159) + %constant.4160 = u32[] constant(0) + ROOT %dynamic-slice.337 = bf16[128,129280]{1,0:T(8,128)(2,1)} dynamic-slice(%all-reduce.144, %multiply.1330, %constant.4160), dynamic_slice_sizes={128,129280}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294966911","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%region_105.125 (psum.0: bf16[], psum.1: bf16[]) -> bf16[] { + %psum.0 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} + %psum.1 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} + ROOT %add.1448 = bf16[]{:T(256)} add(%psum.0, %psum.1), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%all-reduce-scatter.6 (input.6: bf16[256,512,512]) -> bf16[256,128,512] { + %input.6 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %all-reduce.146 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} all-reduce(%input.6), channel_id=1, replica_groups={{0,1,2,3}}, use_global_device_ids=true, to_apply=%region_105.125, frontend_attributes={from-cross-replica-sharding="true"}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"2"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + %constant.4163 = u32[] constant(0) + %partition-id.65 = u32[] partition-id() + %constant.4164 = u32[] constant(128) + %multiply.1332 = u32[]{:T(128)} multiply(%partition-id.65, %constant.4164) + ROOT %dynamic-slice.339 = bf16[256,128,512]{2,1,0:T(8,128)(2,1)} dynamic-slice(%all-reduce.146, %constant.4163, %multiply.1332, %constant.4163), dynamic_slice_sizes={256,128,512}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294966911","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%region_106.126 (psum.4: bf16[], psum.5: bf16[]) -> bf16[] { + %psum.4 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} + %psum.5 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} + ROOT %add.1449 = bf16[]{:T(256)} add(%psum.4, %psum.5), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%all-reduce-scatter.7 (input.7: bf16[256,512,512]) -> bf16[256,128,512] { + %input.7 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %all-reduce.147 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} all-reduce(%input.7), channel_id=1, replica_groups={{0,1,2,3}}, use_global_device_ids=true, to_apply=%region_106.126, frontend_attributes={from-cross-replica-sharding="true"}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"2"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + %constant.4166 = u32[] constant(0) + %partition-id.66 = u32[] partition-id() + %constant.4167 = u32[] constant(128) + %multiply.1333 = u32[]{:T(128)} multiply(%partition-id.66, %constant.4167) + ROOT %dynamic-slice.340 = bf16[256,128,512]{2,1,0:T(8,128)(2,1)} dynamic-slice(%all-reduce.147, %constant.4166, %multiply.1333, %constant.4166), dynamic_slice_sizes={256,128,512}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294966911","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%add.21.clone (x.43: bf16[], y.43: bf16[]) -> bf16[] { + %x.43 = bf16[]{:T(256)} parameter(0) + %y.43 = bf16[]{:T(256)} parameter(1) + ROOT %add.1924 = bf16[]{:T(256)} add(%x.43, %y.43), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%all-reduce-scatter.10 (input.10: bf16[512,128,256]) -> bf16[128,128,256] { + %input.10 = bf16[512,128,256]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %all-reduce.150 = bf16[512,128,256]{2,0,1:T(8,128)(2,1)} all-reduce(%input.10), channel_id=95, replica_groups={{0,1,2,3}}, use_global_device_ids=true, to_apply=%add.21.clone, frontend_attributes={from-cross-replica-sharding="true"}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"2"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + %partition-id.69 = u32[] partition-id() + %constant.4173 = u32[] constant(128) + %multiply.1336 = u32[]{:T(128)} multiply(%partition-id.69, %constant.4173) + %constant.4174 = u32[] constant(0) + ROOT %dynamic-slice.343 = bf16[128,128,256]{2,0,1:T(8,128)(2,1)} dynamic-slice(%all-reduce.150, %multiply.1336, %constant.4174, %constant.4174), dynamic_slice_sizes={128,128,256}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294966911","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%add.23.clone (x.47: bf16[], y.47: bf16[]) -> bf16[] { + %x.47 = bf16[]{:T(256)} parameter(0) + %y.47 = bf16[]{:T(256)} parameter(1) + ROOT %add.1928 = bf16[]{:T(256)} add(%x.47, %y.47), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%all-reduce-scatter.11 (input.11: bf16[1536,128,192]) -> bf16[384,128,192] { + %input.11 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} parameter(0) + %all-reduce.151 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} all-reduce(%input.11), channel_id=99, replica_groups={{0,1,2,3}}, use_global_device_ids=true, to_apply=%add.23.clone, frontend_attributes={from-cross-replica-sharding="true"}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"2"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + %partition-id.70 = u32[] partition-id() + %constant.4176 = u32[] constant(384) + %multiply.1337 = u32[]{:T(128)} multiply(%partition-id.70, %constant.4176) + %constant.4177 = u32[] constant(0) + ROOT %dynamic-slice.344 = bf16[384,128,192]{1,0,2:T(8,128)(2,1)} dynamic-slice(%all-reduce.151, %multiply.1337, %constant.4177, %constant.4177), dynamic_slice_sizes={384,128,192}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294965375","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } %compare.1 (name.4: s32[], name.5: s32[], name.6: f32[], name.7: f32[]) -> pred[] { @@ -511,7 +627,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.294 = 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 +635,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.295 = 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 +643,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.296 = 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,1313 +651,898 @@ 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.297 = 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.4519: s32[256]) -> s32[256] { - %param_0.4519 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2073 = s32[256]{0:T(256)} copy(%param_0.4519), 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.4520: s32[256]) -> s32[256] { - %param_0.4520 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2074.cloned.1 = s32[256]{0:T(256)} call(%param_0.4520), to_apply=%called_computation.13 -}, execution_thread="sparsecore" - -%region_49.59 (scatter-add.14: s32[], scatter-add.15: s32[]) -> s32[] { +%region_49.60 (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.1387 = 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.4521: s32[256], param_1.5326: s32[4096], param_2.4491: s32[4096]) -> s32[256] { - %param_0.4521 = s32[256]{0:T(256)} parameter(0) - %param_1.5326 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3911 = s32[4096]{0:T(1024)} reshape(%param_1.5326), 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.3911), 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.4491 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3912 = s32[4096]{0:T(1024)} reshape(%param_2.4491), 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.3912), 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.4521, %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} +%fused_computation.24.clone.clone (param_0.4888: s32[256], param_1.5478: s32[1024], param_2.4317: s32[1024]) -> s32[256] { + %param_0.4888 = s32[256]{0:T(256)} parameter(0) + %param_1.5478 = s32[1024]{0:T(1024)} parameter(1) + %reshape.4940 = s32[1024]{0:T(1024)} reshape(%param_1.5478), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(clip)/max" stack_frame_id=0} + %transpose.1150 = s32[1024]{0:T(1024)} transpose(%reshape.4940), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(clip)/max" stack_frame_id=0} + %param_2.4317 = s32[1024]{0:T(1024)} parameter(2) + %reshape.4941 = s32[1024]{0:T(1024)} reshape(%param_2.4317), metadata={op_name="jit(train_step)/shard_map/broadcast_in_dim" stack_frame_id=0} + %transpose.1151 = s32[1024]{0:T(1024)} transpose(%reshape.4941), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.231 = s32[256]{0:T(256)} scatter(%param_0.4888, %transpose.1150, %transpose.1151), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_49.60, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.14 (param_0.4522: s32[256], param_1.5327: s32[4096], param_2.4492: s32[4096]) -> s32[256] { - %param_0.4522 = s32[256]{0:T(256)} parameter(0) - %param_1.5327 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4492 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4522, %param_1.5327, %param_2.4492), 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.10 (param_0.4889: s32[256], param_1.5479: s32[1024], param_2.4318: s32[1024]) -> s32[256] { + %param_0.4889 = s32[256]{0:T(256)} parameter(0) + %param_1.5479 = s32[1024]{0:T(1024)} parameter(1) + %param_2.4318 = s32[1024]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.27 = s32[256]{0:T(256)} fusion(%param_0.4889, %param_1.5479, %param_2.4318), kind=kCustom, calls=%fused_computation.24.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["64"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1088","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.4523: s32[256], param_1.5328: s32[4096], param_2.4493: s32[4096]) -> s32[256] { - %param_0.4523 = s32[256]{0:T(256)} parameter(0) - %param_1.5328 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4493 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4523, %param_1.5328, %param_2.4493), 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.10 (param_0.4890: s32[256], param_1.5480: s32[1024], param_2.4319: s32[1024]) -> s32[256] { + %param_0.4890 = s32[256]{0:T(256)} parameter(0) + %param_1.5480 = s32[1024]{0:T(1024)} parameter(1) + %param_2.4319 = s32[1024]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.28.cloned.1 = s32[256]{0:T(256)} call(%param_0.4890, %param_1.5480, %param_2.4319), to_apply=%called_computation.10, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/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.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":[]} - %copy.2074.cloned.1.call-start = ((s32[256]{0:T(256)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%param_0.84), async_execution_thread="sparsecore", calls=%async_computation.13 - %copy.2074.cloned.1.call-done = s32[256]{0:T(256)} async-done(%copy.2074.cloned.1.call-start) - %scatter_offload_custom_fusion.40.cloned.1.call-start = ((s32[256]{0:T(256)}, s32[4096]{0:T(1024)}, s32[4096]{0:T(1024)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%copy.2074.cloned.1.call-done, %param_1.136, %param_2.80), async_execution_thread="sparsecore", calls=%async_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} - 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} +%called_computation (param_0.91: s32[256], param_1.149: s32[1024], param_2.94: s32[1024], param_3.3117: token[]) -> s32[256] { + %param_3.3117 = 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.91 = 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.149 = s32[1024]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_2.94 = s32[1024]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %scatter_offload_custom_fusion.28.cloned.1.call-start = ((s32[256]{0:T(256)}, s32[1024]{0:T(1024)}, s32[1024]{0:T(1024)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%param_0.91, %param_1.149, %param_2.94), async_execution_thread="sparsecore", calls=%async_computation.10, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.28.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.28.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/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) - %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} +%async_computation (param_0.92: s32[256], param_1.150: s32[1024], param_2.95: s32[1024], param_3.3116: token[]) -> s32[256] { + %param_3.3116 = token[] parameter(3) + %param_0.92 = s32[256]{0:T(256)} parameter(0) + %param_1.150 = s32[1024]{0:T(1024)} parameter(1) + %param_2.95 = s32[1024]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.2.cloned.1 = s32[256]{0:T(256)} call(%param_0.92, %param_1.150, %param_2.95, %param_3.3116), to_apply=%called_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.15 (param_0.4524: f32[9]) -> f32[9] { - %param_0.4524 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2075 = f32[9]{0:T(128)} copy(%param_0.4524), 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.11 (param_0.4891: f32[3]) -> f32[3] { + %param_0.4891 = f32[3]{0:T(128)} parameter(0) + ROOT %copy.2090 = f32[3]{0:T(128)} copy(%param_0.4891), 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.4525: f32[9]) -> f32[9] { - %param_0.4525 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2076.cloned.1 = f32[9]{0:T(128)} call(%param_0.4525), to_apply=%called_computation.15 +%async_computation.11 (param_0.4892: f32[3]) -> f32[3] { + %param_0.4892 = f32[3]{0:T(128)} parameter(0) + ROOT %copy.2091.cloned.1 = f32[3]{0:T(128)} call(%param_0.4892), to_apply=%called_computation.11 }, execution_thread="sparsecore" -%region_61.72 (scatter-add.24: f32[], scatter-add.25: f32[]) -> f32[] { +%region_61.73 (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.1393 = 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.4526: f32[9], param_1.5329: s32[256], param_2.4494: f32[256]) -> f32[9] { - %param_0.4526 = f32[9]{0:T(128)} parameter(0) - %param_1.5329 = s32[256]{0:T(256)} parameter(1) - %reshape.3913 = s32[256]{0:T(256)} reshape(%param_1.5329), 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.3913), 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.4494 = f32[256]{0:T(256)} parameter(2) - %reshape.3914 = f32[256]{0:T(256)} reshape(%param_2.4494), 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.3914), 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.4526, %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.4527: f32[9], param_1.5330: s32[256], param_2.4495: f32[256]) -> f32[9] { - %param_0.4527 = f32[9]{0:T(128)} parameter(0) - %param_1.5330 = s32[256]{0:T(256)} parameter(1) - %param_2.4495 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.41 = f32[9]{0:T(128)} fusion(%param_0.4527, %param_1.5330, %param_2.4495), 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.4528: f32[9], param_1.5331: s32[256], param_2.4496: f32[256]) -> f32[9] { - %param_0.4528 = f32[9]{0:T(128)} parameter(0) - %param_1.5331 = s32[256]{0:T(256)} parameter(1) - %param_2.4496 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.42.cloned.1 = f32[9]{0:T(128)} call(%param_0.4528, %param_1.5331, %param_2.4496), 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.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":[]} - %copy.2076.cloned.1.call-start = ((f32[9]{0:T(128)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.87), async_execution_thread="sparsecore", calls=%async_computation.15 - %copy.2076.cloned.1.call-done = f32[9]{0:T(128)} async-done(%copy.2076.cloned.1.call-start) - %scatter_offload_custom_fusion.42.cloned.1.call-start = ((f32[9]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2076.cloned.1.call-done, %param_1.139, %param_2.83), async_execution_thread="sparsecore", calls=%async_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} - 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) - %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} -}, execution_thread="sparsecore" - -%called_computation.17 (param_0.4529: s32[263]) -> s32[263] { - %param_0.4529 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2077 = s32[263]{0:T(512)} copy(%param_0.4529), 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.4530: s32[263]) -> s32[263] { - %param_0.4530 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2078.cloned.1 = s32[263]{0:T(512)} call(%param_0.4530), 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.1394 = 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":[]} +%fused_computation.26.clone.clone (param_0.4893: f32[3], param_1.5481: s32[256], param_2.4320: f32[256]) -> f32[3] { + %param_0.4893 = f32[3]{0:T(128)} parameter(0) + %param_1.5481 = s32[256]{0:T(256)} parameter(1) + %reshape.4942 = s32[256]{0:T(256)} reshape(%param_1.5481), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1152 = s32[256]{0:T(256)} transpose(%reshape.4942), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4320 = f32[256]{0:T(256)} parameter(2) + %reshape.4943 = f32[256]{0:T(256)} reshape(%param_2.4320), metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1153 = f32[256]{0:T(256)} transpose(%reshape.4943), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.232 = f32[3]{0:T(128)} scatter(%param_0.4893, %transpose.1152, %transpose.1153), 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.73, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%fused_computation.25.clone.clone (param_0.4531: s32[263], param_1.5332: s32[8], param_2.4497: s32[8]) -> s32[263] { - %param_0.4531 = s32[263]{0:T(512)} parameter(0) - %param_1.5332 = s32[8]{0:T(128)} parameter(1) - %reshape.3915 = s32[8]{0:T(128)} reshape(%param_1.5332), 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.3915), 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.4497 = s32[8]{0:T(128)} parameter(2) - %reshape.3916 = s32[8]{0:T(128)} reshape(%param_2.4497), 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.3916), 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.4531, %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} +%called_computation.12 (param_0.4894: f32[3], param_1.5482: s32[256], param_2.4321: f32[256]) -> f32[3] { + %param_0.4894 = f32[3]{0:T(128)} parameter(0) + %param_1.5482 = s32[256]{0:T(256)} parameter(1) + %param_2.4321 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.29 = f32[3]{0:T(128)} fusion(%param_0.4894, %param_1.5482, %param_2.4321), kind=kCustom, calls=%fused_computation.26.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/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.18 (param_0.4532: s32[263], param_1.5333: s32[8], param_2.4498: s32[8]) -> s32[263] { - %param_0.4532 = s32[263]{0:T(512)} parameter(0) - %param_1.5333 = s32[8]{0:T(128)} parameter(1) - %param_2.4498 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.43 = s32[263]{0:T(512)} fusion(%param_0.4532, %param_1.5333, %param_2.4498), 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":[]} +%async_computation.12 (param_0.4895: f32[3], param_1.5483: s32[256], param_2.4322: f32[256]) -> f32[3] { + %param_0.4895 = f32[3]{0:T(128)} parameter(0) + %param_1.5483 = s32[256]{0:T(256)} parameter(1) + %param_2.4322 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.30.cloned.1 = f32[3]{0:T(128)} call(%param_0.4895, %param_1.5483, %param_2.4322), to_apply=%called_computation.12, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.18 (param_0.4533: s32[263], param_1.5334: s32[8], param_2.4499: s32[8]) -> s32[263] { - %param_0.4533 = s32[263]{0:T(512)} parameter(0) - %param_1.5334 = s32[8]{0:T(128)} parameter(1) - %param_2.4499 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.44.cloned.1 = s32[263]{0:T(512)} call(%param_0.4533, %param_1.5334, %param_2.4499), 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} +%called_computation.1 (param_0.94: f32[3], param_1.152: s32[256], param_2.97: f32[256], param_3.3129: token[]) -> f32[3] { + %param_3.3129 = 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.94 = f32[3]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_1.152 = 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.97 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %copy.2091.cloned.1.call-start = ((f32[3]{0:T(128)}), f32[3]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.94), async_execution_thread="sparsecore", calls=%async_computation.11 + %copy.2091.cloned.1.call-done = f32[3]{0:T(128)} async-done(%copy.2091.cloned.1.call-start) + %scatter_offload_custom_fusion.30.cloned.1.call-start = ((f32[3]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[3]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2091.cloned.1.call-done, %param_1.152, %param_2.97), async_execution_thread="sparsecore", calls=%async_computation.12, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.30.cloned.1.call-done = f32[3]{0:T(128)} async-done(%scatter_offload_custom_fusion.30.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/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.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":[]} - %copy.2078.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.90), async_execution_thread="sparsecore", calls=%async_computation.17 - %copy.2078.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2078.cloned.1.call-start) - %scatter_offload_custom_fusion.44.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[8]{0:T(128)}, s32[8]{0:T(128)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2078.cloned.1.call-done, %param_1.142, %param_2.86), async_execution_thread="sparsecore", calls=%async_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} - 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} +%async_computation.1 (param_0.95: f32[3], param_1.153: s32[256], param_2.98: f32[256], param_3.3128: token[]) -> f32[3] { + %param_3.3128 = token[] parameter(3) + %param_0.95 = f32[3]{0:T(128)} parameter(0) + %param_1.153 = s32[256]{0:T(256)} parameter(1) + %param_2.98 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.5.cloned.1 = f32[3]{0:T(128)} call(%param_0.95, %param_1.153, %param_2.98, %param_3.3128), to_apply=%called_computation.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/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) - %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} +%called_computation.13 (param_0.4896: s32[257]) -> s32[257] { + %param_0.4896 = s32[257]{0:T(512)} parameter(0) + ROOT %copy.2092 = s32[257]{0:T(512)} copy(%param_0.4896), 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" -%called_computation.19 (param_0.4534: s32[263]) -> s32[263] { - %param_0.4534 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2079 = 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":[]} +%async_computation.13 (param_0.4897: s32[257]) -> s32[257] { + %param_0.4897 = s32[257]{0:T(512)} parameter(0) + ROOT %copy.2093.cloned.1 = s32[257]{0:T(512)} call(%param_0.4897), to_apply=%called_computation.13 }, execution_thread="sparsecore" -%async_computation.19 (param_0.4535: s32[263]) -> s32[263] { - %param_0.4535 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2080.cloned.1 = s32[263]{0:T(512)} call(%param_0.4535), to_apply=%called_computation.19 -}, execution_thread="sparsecore" - -%region_73.86.clone (scatter-add.163: s32[], scatter-add.164: s32[]) -> s32[] { +%region_73.87.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.2527 = 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.2618 = 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.4536: s32[263], param_1.5335: s32[256], param_2.4500: s32[256]) -> s32[263] { - %param_0.4536 = s32[263]{0:T(512)} parameter(0) - %param_1.5335 = s32[256]{0:T(256)} parameter(1) - %reshape.3917 = s32[256]{0:T(256)} reshape(%param_1.5335), metadata={op_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.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(gmm)/select_n" stack_frame_id=0} - %param_2.4500 = s32[256]{0:T(256)} parameter(2) - %reshape.3918 = s32[256]{0:T(256)} reshape(%param_2.4500), 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.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.234 = s32[263]{0:T(512)} scatter(%param_0.4536, %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} +%fused_computation.28.clone.clone (param_0.4898: s32[257], param_1.5484: s32[256], param_2.4323: s32[256]) -> s32[257] { + %param_0.4898 = s32[257]{0:T(512)} parameter(0) + %param_1.5484 = s32[256]{0:T(256)} parameter(1) + %reshape.4944 = s32[256]{0:T(256)} reshape(%param_1.5484), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1154 = s32[256]{0:T(256)} transpose(%reshape.4944), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4323 = s32[256]{0:T(256)} parameter(2) + %reshape.4945 = s32[256]{0:T(256)} reshape(%param_2.4323), metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1155 = s32[256]{0:T(256)} transpose(%reshape.4945), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.233 = s32[257]{0:T(512)} scatter(%param_0.4898, %transpose.1154, %transpose.1155), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_73.87.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.20 (param_0.4537: s32[263], param_1.5336: s32[256], param_2.4501: s32[256]) -> s32[263] { - %param_0.4537 = s32[263]{0:T(512)} parameter(0) - %param_1.5336 = s32[256]{0:T(256)} parameter(1) - %param_2.4501 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.45 = s32[263]{0:T(512)} fusion(%param_0.4537, %param_1.5336, %param_2.4501), 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.14 (param_0.4899: s32[257], param_1.5485: s32[256], param_2.4324: s32[256]) -> s32[257] { + %param_0.4899 = s32[257]{0:T(512)} parameter(0) + %param_1.5485 = s32[256]{0:T(256)} parameter(1) + %param_2.4324 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.31 = s32[257]{0:T(512)} fusion(%param_0.4899, %param_1.5485, %param_2.4324), kind=kCustom, calls=%fused_computation.28.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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.4538: s32[263], param_1.5337: s32[256], param_2.4502: s32[256]) -> s32[263] { - %param_0.4538 = s32[263]{0:T(512)} parameter(0) - %param_1.5337 = s32[256]{0:T(256)} parameter(1) - %param_2.4502 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.46.cloned.1 = s32[263]{0:T(512)} call(%param_0.4538, %param_1.5337, %param_2.4502), 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.14 (param_0.4900: s32[257], param_1.5486: s32[256], param_2.4325: s32[256]) -> s32[257] { + %param_0.4900 = s32[257]{0:T(512)} parameter(0) + %param_1.5486 = s32[256]{0:T(256)} parameter(1) + %param_2.4325 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.32.cloned.1 = s32[257]{0:T(512)} call(%param_0.4900, %param_1.5486, %param_2.4325), to_apply=%called_computation.14, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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.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":[]} - %copy.2080.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.93), async_execution_thread="sparsecore", calls=%async_computation.19 - %copy.2080.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2080.cloned.1.call-start) - %scatter_offload_custom_fusion.46.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2080.cloned.1.call-done, %param_1.145, %param_2.89), async_execution_thread="sparsecore", calls=%async_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} - 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} +%called_computation.2 (param_0.97: s32[257], param_1.155: s32[256], param_2.100: s32[256], param_3.3121: token[]) -> s32[257] { + %param_3.3121 = 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.97 = s32[257]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_1.155 = 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.100 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %copy.2093.cloned.1.call-start = ((s32[257]{0:T(512)}), s32[257]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.97), async_execution_thread="sparsecore", calls=%async_computation.13 + %copy.2093.cloned.1.call-done = s32[257]{0:T(512)} async-done(%copy.2093.cloned.1.call-start) + %scatter_offload_custom_fusion.32.cloned.1.call-start = ((s32[257]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[257]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2093.cloned.1.call-done, %param_1.155, %param_2.100), async_execution_thread="sparsecore", calls=%async_computation.14, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.32.cloned.1.call-done = s32[257]{0:T(512)} async-done(%scatter_offload_custom_fusion.32.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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) - %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} +%async_computation.2 (param_0.98: s32[257], param_1.156: s32[256], param_2.101: s32[256], param_3.3120: token[]) -> s32[257] { + %param_3.3120 = token[] parameter(3) + %param_0.98 = s32[257]{0:T(512)} parameter(0) + %param_1.156 = s32[256]{0:T(256)} parameter(1) + %param_2.101 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.8.cloned.1 = s32[257]{0:T(512)} call(%param_0.98, %param_1.156, %param_2.101, %param_3.3120), to_apply=%called_computation.2, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.21 (param_0.4539: f32[9]) -> f32[9] { - %param_0.4539 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2081 = f32[9]{0:T(128)} copy(%param_0.4539), 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.4901: f32[3]) -> f32[3] { + %param_0.4901 = f32[3]{0:T(128)} parameter(0) + ROOT %copy.2094 = f32[3]{0:T(128)} copy(%param_0.4901), 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.4540: f32[9]) -> f32[9] { - %param_0.4540 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2082.cloned.1 = f32[9]{0:T(128)} call(%param_0.4540), to_apply=%called_computation.21 +%async_computation.15 (param_0.4902: f32[3]) -> f32[3] { + %param_0.4902 = f32[3]{0:T(128)} parameter(0) + ROOT %copy.2095.cloned.1 = f32[3]{0:T(128)} call(%param_0.4902), to_apply=%called_computation.15 }, execution_thread="sparsecore" -%region_79.95.clone (scatter-add.167: f32[], scatter-add.168: f32[]) -> f32[] { +%region_79.96.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.2529 = 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.4541: f32[9], param_1.5338: s32[256], param_2.4503: f32[256]) -> f32[9] { - %param_0.4541 = f32[9]{0:T(128)} parameter(0) - %param_1.5338 = s32[256]{0:T(256)} parameter(1) - %reshape.3919 = s32[256]{0:T(256)} reshape(%param_1.5338), metadata={op_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.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(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4503 = f32[256]{0:T(256)} parameter(2) - %reshape.3920 = f32[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.1109 = f32[256]{0:T(256)} transpose(%reshape.3920), 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.4541, %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.4542: f32[9], param_1.5339: s32[256], param_2.4504: f32[256]) -> f32[9] { - %param_0.4542 = f32[9]{0:T(128)} parameter(0) - %param_1.5339 = s32[256]{0:T(256)} parameter(1) - %param_2.4504 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.47 = f32[9]{0:T(128)} fusion(%param_0.4542, %param_1.5339, %param_2.4504), 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.4543: f32[9], param_1.5340: s32[256], param_2.4505: f32[256]) -> f32[9] { - %param_0.4543 = f32[9]{0:T(128)} parameter(0) - %param_1.5340 = s32[256]{0:T(256)} parameter(1) - %param_2.4505 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.48.cloned.1 = f32[9]{0:T(128)} call(%param_0.4543, %param_1.5340, %param_2.4505), 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.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":[]} - %copy.2082.cloned.1.call-start = ((f32[9]{0:T(128)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.96), async_execution_thread="sparsecore", calls=%async_computation.21 - %copy.2082.cloned.1.call-done = f32[9]{0:T(128)} async-done(%copy.2082.cloned.1.call-start) - %scatter_offload_custom_fusion.48.cloned.1.call-start = ((f32[9]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2082.cloned.1.call-done, %param_1.148, %param_2.92), async_execution_thread="sparsecore", calls=%async_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} - 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) - %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} -}, execution_thread="sparsecore" - -%called_computation.23 (param_0.4544: s32[263]) -> s32[263] { - %param_0.4544 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2083 = s32[263]{0:T(512)} copy(%param_0.4544), 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.4545: s32[263]) -> s32[263] { - %param_0.4545 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2084.cloned.1 = s32[263]{0:T(512)} call(%param_0.4545), 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.2531 = 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.2620 = 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.28.clone.clone (param_0.4546: s32[263], param_1.5341: s32[8], param_2.4506: s32[8]) -> s32[263] { - %param_0.4546 = s32[263]{0:T(512)} parameter(0) - %param_1.5341 = s32[8]{0:T(128)} parameter(1) - %reshape.3921 = s32[8]{0:T(128)} reshape(%param_1.5341), metadata={op_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.3921), 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.4506 = s32[8]{0:T(128)} parameter(2) - %reshape.3922 = s32[8]{0:T(128)} reshape(%param_2.4506), 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.3922), 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.4546, %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} +%fused_computation.29.clone.clone (param_0.4903: f32[3], param_1.5487: s32[256], param_2.4326: f32[256]) -> f32[3] { + %param_0.4903 = f32[3]{0:T(128)} parameter(0) + %param_1.5487 = s32[256]{0:T(256)} parameter(1) + %reshape.4946 = s32[256]{0:T(256)} reshape(%param_1.5487), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1156 = s32[256]{0:T(256)} transpose(%reshape.4946), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4326 = f32[256]{0:T(256)} parameter(2) + %reshape.4947 = f32[256]{0:T(256)} reshape(%param_2.4326), metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1157 = f32[256]{0:T(256)} transpose(%reshape.4947), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.234 = f32[3]{0:T(128)} scatter(%param_0.4903, %transpose.1156, %transpose.1157), 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.96.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.24 (param_0.4547: s32[263], param_1.5342: s32[8], param_2.4507: s32[8]) -> s32[263] { - %param_0.4547 = s32[263]{0:T(512)} parameter(0) - %param_1.5342 = s32[8]{0:T(128)} parameter(1) - %param_2.4507 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.49 = s32[263]{0:T(512)} fusion(%param_0.4547, %param_1.5342, %param_2.4507), 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.16 (param_0.4904: f32[3], param_1.5488: s32[256], param_2.4327: f32[256]) -> f32[3] { + %param_0.4904 = f32[3]{0:T(128)} parameter(0) + %param_1.5488 = s32[256]{0:T(256)} parameter(1) + %param_2.4327 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.33 = f32[3]{0:T(128)} fusion(%param_0.4904, %param_1.5488, %param_2.4327), kind=kCustom, calls=%fused_computation.29.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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.24 (param_0.4548: s32[263], param_1.5343: s32[8], param_2.4508: s32[8]) -> s32[263] { - %param_0.4548 = s32[263]{0:T(512)} parameter(0) - %param_1.5343 = s32[8]{0:T(128)} parameter(1) - %param_2.4508 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.50.cloned.1 = s32[263]{0:T(512)} call(%param_0.4548, %param_1.5343, %param_2.4508), 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.16 (param_0.4905: f32[3], param_1.5489: s32[256], param_2.4328: f32[256]) -> f32[3] { + %param_0.4905 = f32[3]{0:T(128)} parameter(0) + %param_1.5489 = s32[256]{0:T(256)} parameter(1) + %param_2.4328 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.34.cloned.1 = f32[3]{0:T(128)} call(%param_0.4905, %param_1.5489, %param_2.4328), to_apply=%called_computation.16, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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.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":[]} - %copy.2084.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.99), async_execution_thread="sparsecore", calls=%async_computation.23 - %copy.2084.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2084.cloned.1.call-start) - %scatter_offload_custom_fusion.50.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[8]{0:T(128)}, s32[8]{0:T(128)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2084.cloned.1.call-done, %param_1.151, %param_2.95), async_execution_thread="sparsecore", calls=%async_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} - 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} +%called_computation.3 (param_0.100: f32[3], param_1.158: s32[256], param_2.103: f32[256], param_3.3127: token[]) -> f32[3] { + %param_3.3127 = 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.100 = f32[3]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_1.158 = 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.103 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %copy.2095.cloned.1.call-start = ((f32[3]{0:T(128)}), f32[3]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.100), async_execution_thread="sparsecore", calls=%async_computation.15 + %copy.2095.cloned.1.call-done = f32[3]{0:T(128)} async-done(%copy.2095.cloned.1.call-start) + %scatter_offload_custom_fusion.34.cloned.1.call-start = ((f32[3]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[3]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2095.cloned.1.call-done, %param_1.158, %param_2.103), async_execution_thread="sparsecore", calls=%async_computation.16, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.34.cloned.1.call-done = f32[3]{0:T(128)} async-done(%scatter_offload_custom_fusion.34.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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) - %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} +%async_computation.3 (param_0.101: f32[3], param_1.159: s32[256], param_2.104: f32[256], param_3.3126: token[]) -> f32[3] { + %param_3.3126 = token[] parameter(3) + %param_0.101 = f32[3]{0:T(128)} parameter(0) + %param_1.159 = s32[256]{0:T(256)} parameter(1) + %param_2.104 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.11.cloned.1 = f32[3]{0:T(128)} call(%param_0.101, %param_1.159, %param_2.104, %param_3.3126), to_apply=%called_computation.3, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.25 (param_0.4549: s32[263]) -> s32[263] { - %param_0.4549 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2085 = 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.25 (param_0.4550: s32[263]) -> s32[263] { - %param_0.4550 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2086.cloned.1 = s32[263]{0:T(512)} call(%param_0.4550), to_apply=%called_computation.25 -}, execution_thread="sparsecore" - -%region_96.114 (scatter-add.48: s32[], scatter-add.49: s32[]) -> s32[] { +%region_94.113 (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.1434 = 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.1436 = 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.4551: s32[263], param_1.5344: s32[256], param_2.4509: s32[256]) -> s32[263] { - %param_0.4551 = s32[263]{0:T(512)} parameter(0) - %param_1.5344 = s32[256]{0:T(256)} parameter(1) - %reshape.3923 = s32[256]{0:T(256)} reshape(%param_1.5344), metadata={op_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.3923), 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.4509 = s32[256]{0:T(256)} parameter(2) - %reshape.3924 = s32[256]{0:T(256)} reshape(%param_2.4509), 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.3924), 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.4551, %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} +%fused_computation.31.clone.clone (param_0.4908: s32[257], param_1.5490: s32[256], param_2.4329: s32[256]) -> s32[257] { + %param_0.4908 = s32[257]{0:T(512)} parameter(0) + %param_1.5490 = s32[256]{0:T(256)} parameter(1) + %reshape.4948 = s32[256]{0:T(256)} reshape(%param_1.5490), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %transpose.1158 = s32[256]{0:T(256)} transpose(%reshape.4948), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/select_n" stack_frame_id=0} + %param_2.4329 = s32[256]{0:T(256)} parameter(2) + %reshape.4949 = s32[256]{0:T(256)} reshape(%param_2.4329), metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1159 = s32[256]{0:T(256)} transpose(%reshape.4949), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.235 = s32[257]{0:T(512)} scatter(%param_0.4908, %transpose.1158, %transpose.1159), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_94.113, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.26 (param_0.4552: s32[263], param_1.5345: s32[256], param_2.4510: s32[256]) -> s32[263] { - %param_0.4552 = s32[263]{0:T(512)} parameter(0) - %param_1.5345 = s32[256]{0:T(256)} parameter(1) - %param_2.4510 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.51 = s32[263]{0:T(512)} fusion(%param_0.4552, %param_1.5345, %param_2.4510), 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.18 (param_0.4909: s32[257], param_1.5491: s32[256], param_2.4330: s32[256]) -> s32[257] { + %param_0.4909 = s32[257]{0:T(512)} parameter(0) + %param_1.5491 = s32[256]{0:T(256)} parameter(1) + %param_2.4330 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.35 = s32[257]{0:T(512)} fusion(%param_0.4909, %param_1.5491, %param_2.4330), kind=kCustom, calls=%fused_computation.31.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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.4553: s32[263], param_1.5346: s32[256], param_2.4511: s32[256]) -> s32[263] { - %param_0.4553 = s32[263]{0:T(512)} parameter(0) - %param_1.5346 = s32[256]{0:T(256)} parameter(1) - %param_2.4511 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.52.cloned.1 = s32[263]{0:T(512)} call(%param_0.4553, %param_1.5346, %param_2.4511), 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.18 (param_0.4910: s32[257], param_1.5492: s32[256], param_2.4331: s32[256]) -> s32[257] { + %param_0.4910 = s32[257]{0:T(512)} parameter(0) + %param_1.5492 = s32[256]{0:T(256)} parameter(1) + %param_2.4331 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.36.cloned.1 = s32[257]{0:T(512)} call(%param_0.4910, %param_1.5492, %param_2.4331), to_apply=%called_computation.18, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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.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":[]} - %copy.2086.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.102), async_execution_thread="sparsecore", calls=%async_computation.25 - %copy.2086.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2086.cloned.1.call-start) - %scatter_offload_custom_fusion.52.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2086.cloned.1.call-done, %param_1.154, %param_2.98), async_execution_thread="sparsecore", calls=%async_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} - 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} +%called_computation.4 (param_0.103: s32[257], param_1.161: s32[256], param_2.106: s32[256], param_3.3123: token[]) -> s32[257] { + %param_3.3123 = 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.103 = s32[257]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_1.161 = 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.106 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %scatter_offload_custom_fusion.36.cloned.1.call-start = ((s32[257]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[257]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.103, %param_1.161, %param_2.106), async_execution_thread="sparsecore", calls=%async_computation.18, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.36.cloned.1.call-done = s32[257]{0:T(512)} async-done(%scatter_offload_custom_fusion.36.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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) - %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} +%async_computation.4 (param_0.104: s32[257], param_1.162: s32[256], param_2.107: s32[256], param_3.3122: token[]) -> s32[257] { + %param_3.3122 = token[] parameter(3) + %param_0.104 = s32[257]{0:T(512)} parameter(0) + %param_1.162 = s32[256]{0:T(256)} parameter(1) + %param_2.107 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.14.cloned.1 = s32[257]{0:T(512)} call(%param_0.104, %param_1.162, %param_2.107, %param_3.3122), to_apply=%called_computation.4, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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[] { +%region_100.119 (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.1437 = 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.4556: f32[9], param_1.5347: s32[256], param_2.4512: f32[256]) -> f32[9] { - %param_0.4556 = f32[9]{0:T(128)} parameter(0) - %param_1.5347 = s32[256]{0:T(256)} parameter(1) - %reshape.3925 = s32[256]{0:T(256)} reshape(%param_1.5347), metadata={op_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.3925), 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.4512 = f32[256]{0:T(256)} parameter(2) - %reshape.3926 = f32[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.1115 = 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.4556, %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.4557: f32[9], param_1.5348: s32[256], param_2.4513: f32[256]) -> f32[9] { - %param_0.4557 = f32[9]{0:T(128)} parameter(0) - %param_1.5348 = s32[256]{0:T(256)} parameter(1) - %param_2.4513 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.53 = f32[9]{0:T(128)} fusion(%param_0.4557, %param_1.5348, %param_2.4513), 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.4558: f32[9], param_1.5349: s32[256], param_2.4514: f32[256]) -> f32[9] { - %param_0.4558 = f32[9]{0:T(128)} parameter(0) - %param_1.5349 = s32[256]{0:T(256)} parameter(1) - %param_2.4514 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.54.cloned.1 = f32[9]{0:T(128)} call(%param_0.4558, %param_1.5349, %param_2.4514), 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.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":[]} - %scatter_offload_custom_fusion.54.cloned.1.call-start = ((f32[9]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.105, %param_1.157, %param_2.101), async_execution_thread="sparsecore", calls=%async_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} - 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} + ROOT %add.1439 = 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" -%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.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} +%fused_computation.32.clone.clone (param_0.4913: f32[3], param_1.5493: s32[256], param_2.4332: f32[256]) -> f32[3] { + %param_0.4913 = f32[3]{0:T(128)} parameter(0) + %param_1.5493 = s32[256]{0:T(256)} parameter(1) + %reshape.4950 = s32[256]{0:T(256)} reshape(%param_1.5493), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1160 = s32[256]{0:T(256)} transpose(%reshape.4950), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4332 = f32[256]{0:T(256)} parameter(2) + %reshape.4951 = f32[256]{0:T(256)} reshape(%param_2.4332), metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1161 = f32[256]{0:T(256)} transpose(%reshape.4951), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.236 = f32[3]{0:T(128)} scatter(%param_0.4913, %transpose.1160, %transpose.1161), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_100.119, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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.1438 = 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":[]} +%called_computation.20 (param_0.4914: f32[3], param_1.5494: s32[256], param_2.4333: f32[256]) -> f32[3] { + %param_0.4914 = f32[3]{0:T(128)} parameter(0) + %param_1.5494 = s32[256]{0:T(256)} parameter(1) + %param_2.4333 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.37 = f32[3]{0:T(128)} fusion(%param_0.4914, %param_1.5494, %param_2.4333), kind=kCustom, calls=%fused_computation.32.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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" -%fused_computation.31.clone.clone (param_0.4561: s32[263], param_1.5350: s32[8], param_2.4515: s32[8]) -> s32[263] { - %param_0.4561 = s32[263]{0:T(512)} parameter(0) - %param_1.5350 = s32[8]{0:T(128)} parameter(1) - %reshape.3927 = s32[8]{0:T(128)} reshape(%param_1.5350), metadata={op_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.3927), 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.4515 = s32[8]{0:T(128)} parameter(2) - %reshape.3928 = s32[8]{0:T(128)} reshape(%param_2.4515), 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.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.4561, %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} +%async_computation.20 (param_0.4915: f32[3], param_1.5495: s32[256], param_2.4334: f32[256]) -> f32[3] { + %param_0.4915 = f32[3]{0:T(128)} parameter(0) + %param_1.5495 = s32[256]{0:T(256)} parameter(1) + %param_2.4334 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.38.cloned.1 = f32[3]{0:T(128)} call(%param_0.4915, %param_1.5495, %param_2.4334), to_apply=%called_computation.20, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.30 (param_0.4562: s32[263], param_1.5351: s32[8], param_2.4516: s32[8]) -> s32[263] { - %param_0.4562 = s32[263]{0:T(512)} parameter(0) - %param_1.5351 = s32[8]{0:T(128)} parameter(1) - %param_2.4516 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.55 = s32[263]{0:T(512)} fusion(%param_0.4562, %param_1.5351, %param_2.4516), 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.5 (param_0.106: f32[3], param_1.164: s32[256], param_2.109: f32[256], param_3.3131: token[]) -> f32[3] { + %param_3.3131 = 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.106 = f32[3]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_1.164 = 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.109 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %scatter_offload_custom_fusion.38.cloned.1.call-start = ((f32[3]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[3]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.106, %param_1.164, %param_2.109), async_execution_thread="sparsecore", calls=%async_computation.20, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.38.cloned.1.call-done = f32[3]{0:T(128)} async-done(%scatter_offload_custom_fusion.38.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%async_computation.30 (param_0.4563: s32[263], param_1.5352: s32[8], param_2.4517: s32[8]) -> s32[263] { - %param_0.4563 = s32[263]{0:T(512)} parameter(0) - %param_1.5352 = s32[8]{0:T(128)} parameter(1) - %param_2.4517 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.56.cloned.1 = s32[263]{0:T(512)} call(%param_0.4563, %param_1.5352, %param_2.4517), 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.5 (param_0.107: f32[3], param_1.165: s32[256], param_2.110: f32[256], param_3.3130: token[]) -> f32[3] { + %param_3.3130 = token[] parameter(3) + %param_0.107 = f32[3]{0:T(128)} parameter(0) + %param_1.165 = s32[256]{0:T(256)} parameter(1) + %param_2.110 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.17.cloned.1 = f32[3]{0:T(128)} call(%param_0.107, %param_1.165, %param_2.110, %param_3.3130), to_apply=%called_computation.5, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/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":[]} - %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":[]} - %scatter_offload_custom_fusion.56.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[8]{0:T(128)}, s32[8]{0:T(128)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.108, %param_1.160, %param_2.104), async_execution_thread="sparsecore", calls=%async_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} - 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} +%called_computation.21 (param_0.4916: s32[256]) -> s32[256] { + %param_0.4916 = s32[256]{0:T(256)} parameter(0) + ROOT %copy.2100 = s32[256]{0:T(256)} copy(%param_0.4916), 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.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.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} +%async_computation.21 (param_0.4917: s32[256]) -> s32[256] { + %param_0.4917 = s32[256]{0:T(256)} parameter(0) + ROOT %copy.2101.cloned.1 = s32[256]{0:T(256)} call(%param_0.4917), to_apply=%called_computation.21 }, 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.1345 = 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.1343 = 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.4566: s32[256], param_1.5353: s32[4096], param_2.4518: s32[4096]) -> s32[256] { - %param_0.4566 = s32[256]{0:T(256)} parameter(0) - %param_1.5353 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3929 = s32[4096]{0:T(1024)} reshape(%param_1.5353), 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.3929), 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.4518 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3930 = s32[4096]{0:T(1024)} reshape(%param_2.4518), 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.3930), 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.4566, %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} +%fused_computation.19.clone.clone.clone (param_0.4918: s32[256], param_1.5496: s32[1024], param_2.4335: s32[1024]) -> s32[256] { + %param_0.4918 = s32[256]{0:T(256)} parameter(0) + %param_1.5496 = s32[1024]{0:T(1024)} parameter(1) + %reshape.4952 = s32[1024]{0:T(1024)} reshape(%param_1.5496), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/select_n" stack_frame_id=0} + %transpose.1162 = s32[1024]{0:T(1024)} transpose(%reshape.4952), dimensions={0}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/select_n" stack_frame_id=0} + %param_2.4335 = s32[1024]{0:T(1024)} parameter(2) + %reshape.4953 = s32[1024]{0:T(1024)} reshape(%param_2.4335), metadata={op_name="jit(train_step)/shard_map/broadcast_in_dim" stack_frame_id=0} + %transpose.1163 = s32[1024]{0:T(1024)} transpose(%reshape.4953), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/broadcast_in_dim" stack_frame_id=0} + ROOT %scatter-add.237 = s32[256]{0:T(256)} scatter(%param_0.4918, %transpose.1162, %transpose.1163), 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()/while/body/closed_call/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.32 (param_0.4567: s32[256], param_1.5354: s32[4096], param_2.4519: s32[4096]) -> s32[256] { - %param_0.4567 = s32[256]{0:T(256)} parameter(0) - %param_1.5354 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4519 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.57 = s32[256]{0:T(256)} fusion(%param_0.4567, %param_1.5354, %param_2.4519), 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.22 (param_0.4919: s32[256], param_1.5497: s32[1024], param_2.4336: s32[1024]) -> s32[256] { + %param_0.4919 = s32[256]{0:T(256)} parameter(0) + %param_1.5497 = s32[1024]{0:T(1024)} parameter(1) + %param_2.4336 = s32[1024]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4919, %param_1.5497, %param_2.4336), kind=kCustom, calls=%fused_computation.19.clone.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["64"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1088","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.4568: s32[256], param_1.5355: s32[4096], param_2.4520: s32[4096]) -> s32[256] { - %param_0.4568 = s32[256]{0:T(256)} parameter(0) - %param_1.5355 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4520 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.58.cloned.1 = s32[256]{0:T(256)} call(%param_0.4568, %param_1.5355, %param_2.4520), 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.22 (param_0.4920: s32[256], param_1.5498: s32[1024], param_2.4337: s32[1024]) -> s32[256] { + %param_0.4920 = s32[256]{0:T(256)} parameter(0) + %param_1.5498 = s32[1024]{0:T(1024)} parameter(1) + %param_2.4337 = s32[1024]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4920, %param_1.5498, %param_2.4337), to_apply=%called_computation.22, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/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.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":[]} - %scatter_offload_custom_fusion.58.cloned.1.call-start = ((s32[256]{0:T(256)}, s32[4096]{0:T(1024)}, s32[4096]{0:T(1024)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%param_0.111, %param_1.163, %param_2.107), async_execution_thread="sparsecore", calls=%async_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} - 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} +%called_computation.6 (param_0.109: s32[256], param_1.167: s32[1024], param_2.112: s32[1024], param_3.3115: token[]) -> s32[256] { + %param_3.3115 = 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.109 = 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.167 = s32[1024]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_2.112 = s32[1024]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %copy.2101.cloned.1.call-start = ((s32[256]{0:T(256)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%param_0.109), async_execution_thread="sparsecore", calls=%async_computation.21 + %copy.2101.cloned.1.call-done = s32[256]{0:T(256)} async-done(%copy.2101.cloned.1.call-start) + %scatter_offload_custom_fusion.40.cloned.1.call-start = ((s32[256]{0:T(256)}, s32[1024]{0:T(1024)}, s32[1024]{0:T(1024)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%copy.2101.cloned.1.call-done, %param_1.167, %param_2.112), async_execution_thread="sparsecore", calls=%async_computation.22, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/scatter-add" stack_frame_id=0} + 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)/jvp()/while/body/closed_call/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) - %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} +%async_computation.6 (param_0.110: s32[256], param_1.168: s32[1024], param_2.113: s32[1024], param_3.3114: token[]) -> s32[256] { + %param_3.3114 = token[] parameter(3) + %param_0.110 = s32[256]{0:T(256)} parameter(0) + %param_1.168 = s32[1024]{0:T(1024)} parameter(1) + %param_2.113 = s32[1024]{0:T(1024)} parameter(2) + ROOT %scatter_offload_custom_fusion.20.cloned.1 = s32[256]{0:T(256)} call(%param_0.110, %param_1.168, %param_2.113, %param_3.3114), to_apply=%called_computation.6, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.33 (param_0.4569: s32[263]) -> s32[263] { - %param_0.4569 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2093 = s32[263]{0:T(512)} copy(%param_0.4569), 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.4921: s32[257]) -> s32[257] { + %param_0.4921 = s32[257]{0:T(512)} parameter(0) + ROOT %copy.2102 = s32[257]{0:T(512)} copy(%param_0.4921), 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.4570: s32[263]) -> s32[263] { - %param_0.4570 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2094.cloned.1 = s32[263]{0:T(512)} call(%param_0.4570), to_apply=%called_computation.33 +%async_computation.23 (param_0.4922: s32[257]) -> s32[257] { + %param_0.4922 = s32[257]{0:T(512)} parameter(0) + ROOT %copy.2103.cloned.1 = s32[257]{0:T(512)} call(%param_0.4922), to_apply=%called_computation.23 }, 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.2516 = 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.2607 = 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.4571: s32[263], param_1.5356: s32[256], param_2.4521: s32[256]) -> s32[263] { - %param_0.4571 = s32[263]{0:T(512)} parameter(0) - %param_1.5356 = s32[256]{0:T(256)} parameter(1) - %reshape.3931 = s32[256]{0:T(256)} reshape(%param_1.5356), 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.3931), 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.4521 = s32[256]{0:T(256)} parameter(2) - %reshape.3932 = s32[256]{0:T(256)} reshape(%param_2.4521) - %transpose.1121 = s32[256]{0:T(256)} transpose(%reshape.3932), dimensions={0} - ROOT %scatter-add.241 = s32[263]{0:T(512)} scatter(%param_0.4571, %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} +%fused_computation.20.clone.clone.clone (param_0.4923: s32[257], param_1.5499: s32[256], param_2.4338: s32[256]) -> s32[257] { + %param_0.4923 = s32[257]{0:T(512)} parameter(0) + %param_1.5499 = s32[256]{0:T(256)} parameter(1) + %reshape.4954 = s32[256]{0:T(256)} reshape(%param_1.5499), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %transpose.1164 = s32[256]{0:T(256)} transpose(%reshape.4954), dimensions={0}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/select_n" stack_frame_id=0} + %param_2.4338 = s32[256]{0:T(256)} parameter(2) + %reshape.4955 = s32[256]{0:T(256)} reshape(%param_2.4338) + %transpose.1165 = s32[256]{0:T(256)} transpose(%reshape.4955), dimensions={0} + ROOT %scatter-add.238 = s32[257]{0:T(512)} scatter(%param_0.4923, %transpose.1164, %transpose.1165), 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()/while/body/closed_call/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.34 (param_0.4572: s32[263], param_1.5357: s32[256], param_2.4522: s32[256]) -> s32[263] { - %param_0.4572 = s32[263]{0:T(512)} parameter(0) - %param_1.5357 = s32[256]{0:T(256)} parameter(1) - %param_2.4522 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.59 = s32[263]{0:T(512)} fusion(%param_0.4572, %param_1.5357, %param_2.4522), 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.24 (param_0.4924: s32[257], param_1.5500: s32[256], param_2.4339: s32[256]) -> s32[257] { + %param_0.4924 = s32[257]{0:T(512)} parameter(0) + %param_1.5500 = s32[256]{0:T(256)} parameter(1) + %param_2.4339 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.41 = s32[257]{0:T(512)} fusion(%param_0.4924, %param_1.5500, %param_2.4339), kind=kCustom, calls=%fused_computation.20.clone.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/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.4573: s32[263], param_1.5358: s32[256], param_2.4523: s32[256]) -> s32[263] { - %param_0.4573 = s32[263]{0:T(512)} parameter(0) - %param_1.5358 = s32[256]{0:T(256)} parameter(1) - %param_2.4523 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.60.cloned.1 = s32[263]{0:T(512)} call(%param_0.4573, %param_1.5358, %param_2.4523), 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.24 (param_0.4925: s32[257], param_1.5501: s32[256], param_2.4340: s32[256]) -> s32[257] { + %param_0.4925 = s32[257]{0:T(512)} parameter(0) + %param_1.5501 = s32[256]{0:T(256)} parameter(1) + %param_2.4340 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.42.cloned.1 = s32[257]{0:T(512)} call(%param_0.4925, %param_1.5501, %param_2.4340), to_apply=%called_computation.24, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/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.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":[]} - %copy.2094.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.114), async_execution_thread="sparsecore", calls=%async_computation.33 - %copy.2094.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2094.cloned.1.call-start) - %scatter_offload_custom_fusion.60.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2094.cloned.1.call-done, %param_1.166, %param_2.110), async_execution_thread="sparsecore", calls=%async_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} - 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} +%called_computation.7 (param_0.112: s32[257], param_1.170: s32[256], param_2.115: s32[256], param_3.3119: token[]) -> s32[257] { + %param_3.3119 = 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.112 = s32[257]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_1.170 = 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.115 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %copy.2103.cloned.1.call-start = ((s32[257]{0:T(512)}), s32[257]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.112), async_execution_thread="sparsecore", calls=%async_computation.23 + %copy.2103.cloned.1.call-done = s32[257]{0:T(512)} async-done(%copy.2103.cloned.1.call-start) + %scatter_offload_custom_fusion.42.cloned.1.call-start = ((s32[257]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[257]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2103.cloned.1.call-done, %param_1.170, %param_2.115), async_execution_thread="sparsecore", calls=%async_computation.24, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.42.cloned.1.call-done = s32[257]{0:T(512)} async-done(%scatter_offload_custom_fusion.42.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/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) - %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} +%async_computation.7 (param_0.113: s32[257], param_1.171: s32[256], param_2.116: s32[256], param_3.3118: token[]) -> s32[257] { + %param_3.3118 = token[] parameter(3) + %param_0.113 = s32[257]{0:T(512)} parameter(0) + %param_1.171 = s32[256]{0:T(256)} parameter(1) + %param_2.116 = s32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.23.cloned.1 = s32[257]{0:T(512)} call(%param_0.113, %param_1.171, %param_2.116, %param_3.3118), to_apply=%called_computation.7, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} }, execution_thread="sparsecore" -%called_computation.35 (param_0.4574: f32[9]) -> f32[9] { - %param_0.4574 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2095 = f32[9]{0:T(128)} copy(%param_0.4574), 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.25 (param_0.4926: f32[3]) -> f32[3] { + %param_0.4926 = f32[3]{0:T(128)} parameter(0) + ROOT %copy.2104 = f32[3]{0:T(128)} copy(%param_0.4926), 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.4575: f32[9]) -> f32[9] { - %param_0.4575 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2096.cloned.1 = f32[9]{0:T(128)} call(%param_0.4575), to_apply=%called_computation.35 +%async_computation.25 (param_0.4927: f32[3]) -> f32[3] { + %param_0.4927 = f32[3]{0:T(128)} parameter(0) + ROOT %copy.2105.cloned.1 = f32[3]{0:T(128)} call(%param_0.4927), to_apply=%called_computation.25 }, 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.2518 = 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.4576: f32[9], param_1.5359: s32[256], param_2.4524: f32[256]) -> f32[9] { - %param_0.4576 = f32[9]{0:T(128)} parameter(0) - %param_1.5359 = s32[256]{0:T(256)} parameter(1) - %reshape.3933 = s32[256]{0:T(256)} reshape(%param_1.5359), 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.3933), 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.4524 = f32[256]{0:T(256)} parameter(2) - %reshape.3934 = f32[256]{0:T(256)} reshape(%param_2.4524) - %transpose.1123 = f32[256]{0:T(256)} transpose(%reshape.3934), dimensions={0} - ROOT %scatter-add.242 = f32[9]{0:T(128)} scatter(%param_0.4576, %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.4577: f32[9], param_1.5360: s32[256], param_2.4525: f32[256]) -> f32[9] { - %param_0.4577 = f32[9]{0:T(128)} parameter(0) - %param_1.5360 = s32[256]{0:T(256)} parameter(1) - %param_2.4525 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.61 = f32[9]{0:T(128)} fusion(%param_0.4577, %param_1.5360, %param_2.4525), 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.4578: f32[9], param_1.5361: s32[256], param_2.4526: f32[256]) -> f32[9] { - %param_0.4578 = f32[9]{0:T(128)} parameter(0) - %param_1.5361 = s32[256]{0:T(256)} parameter(1) - %param_2.4526 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.62.cloned.1 = f32[9]{0:T(128)} call(%param_0.4578, %param_1.5361, %param_2.4526), 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.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":[]} - %copy.2096.cloned.1.call-start = ((f32[9]{0:T(128)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.117), async_execution_thread="sparsecore", calls=%async_computation.35 - %copy.2096.cloned.1.call-done = f32[9]{0:T(128)} async-done(%copy.2096.cloned.1.call-start) - %scatter_offload_custom_fusion.62.cloned.1.call-start = ((f32[9]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2096.cloned.1.call-done, %param_1.169, %param_2.113), async_execution_thread="sparsecore", calls=%async_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} - 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) - %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} -}, execution_thread="sparsecore" - -%called_computation.37 (param_0.4579: s32[263]) -> s32[263] { - %param_0.4579 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2097 = s32[263]{0:T(512)} copy(%param_0.4579), 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.4580: s32[263]) -> s32[263] { - %param_0.4580 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2098.cloned.1 = s32[263]{0:T(512)} call(%param_0.4580), 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.2520 = 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.4581: s32[263], param_1.5362: s32[8], param_2.4527: s32[8]) -> s32[263] { - %param_0.4581 = s32[263]{0:T(512)} parameter(0) - %param_1.5362 = s32[8]{0:T(128)} parameter(1) - %reshape.3935 = s32[8]{0:T(128)} reshape(%param_1.5362), 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.3935), 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.4527 = s32[8]{0:T(128)} parameter(2) - %reshape.3936 = s32[8]{0:T(128)} reshape(%param_2.4527) - %transpose.1125 = s32[8]{0:T(128)} transpose(%reshape.3936), dimensions={0} - ROOT %scatter-add.243 = s32[263]{0:T(512)} scatter(%param_0.4581, %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.4582: s32[263], param_1.5363: s32[8], param_2.4528: s32[8]) -> s32[263] { - %param_0.4582 = s32[263]{0:T(512)} parameter(0) - %param_1.5363 = s32[8]{0:T(128)} parameter(1) - %param_2.4528 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.63 = s32[263]{0:T(512)} fusion(%param_0.4582, %param_1.5363, %param_2.4528), 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.4583: s32[263], param_1.5364: s32[8], param_2.4529: s32[8]) -> s32[263] { - %param_0.4583 = s32[263]{0:T(512)} parameter(0) - %param_1.5364 = s32[8]{0:T(128)} parameter(1) - %param_2.4529 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.64.cloned.1 = s32[263]{0:T(512)} call(%param_0.4583, %param_1.5364, %param_2.4529), 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.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":[]} - %copy.2098.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.120), async_execution_thread="sparsecore", calls=%async_computation.37 - %copy.2098.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2098.cloned.1.call-start) - %scatter_offload_custom_fusion.64.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[8]{0:T(128)}, s32[8]{0:T(128)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2098.cloned.1.call-done, %param_1.172, %param_2.116), async_execution_thread="sparsecore", calls=%async_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} - 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) - %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} -}, execution_thread="sparsecore" - -%region_154.179 (reduce_sum.431: f32[], reduce_sum.254: f32[]) -> f32[] { - %reduce_sum.431 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.254 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4166: f32[3,1536,128,192]) -> f32[] { - %param_0.4166 = 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.4166), 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.5101 = f32[]{:T(128)} constant(0) - ROOT %reduce.669 = f32[]{:T(128)} reduce(%square.564, %constant.5101), 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.1420: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { - %param_0.1420 = 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.1420), 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[] { - %reduce_sum.893 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.603 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.604 = f32[]{:T(128)} add(%reduce_sum.893, %reduce_sum.603), 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":[]}} -} - -%region_187.212 (reduce_sum.655: f32[], reduce_sum.449: f32[]) -> f32[] { - %reduce_sum.655 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.449 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4136: f32[1536,3,128,192], param_1.5026: f32[], param_2.4295: 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.4136 = 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.4713.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.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} - %param_5.2006 = f32[]{:T(128)} parameter(5) - %div.2562.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.2561.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.1374.clone.1, %div.2562.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.2561.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4860.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4272.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="broadcast.334"} - %mul.4719.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %broadcast.4272.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.4864.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4720.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)/mul" stack_frame_id=0} - %mul.4718.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.4720.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3488.clone.1 = f32[1536,3,128,192]{2,3,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_2.4295 = f32[]{:T(128)S(6)} parameter(2) - %div.2558.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4295), 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.4863.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4717.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4863.clone.1), 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)} multiply(%integer_pow.399.clone.1, %mul.4717.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.4862.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4716.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4862.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4714.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.4716.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3487.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4715.clone.1, %mul.4714.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5026 = f32[]{:T(128)S(6)} parameter(1) - %div.2557.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5026), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2556.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3487.clone.1, %div.2557.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.2556.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4861.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3486.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4861.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3485.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.157.clone.1, %add.3486.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.2558.clone.1, %add.3485.clone.1), metadata={op_name="multiply.290"} - %div.2555.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3488.clone.1, %multiply.1293.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4712.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4136, %broadcast.4272.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3484.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2555.clone.1, %mul.4712.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4711.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.4713.clone.1, %add.3484.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3483.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4136, %mul.4711.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.3483.clone.1, %add.3483.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5071 = f32[]{:T(128)} constant(0) - %reduce.670 = f32[]{:T(128)} reduce(%square.565, %constant.5071), 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.5071), 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.3483.clone.1, %add.3487.clone.1, %add.3488.clone.1, %reduce.671.clone.1) -} - -%region_160.185 (reduce_sum.473: f32[], reduce_sum.293: f32[]) -> f32[] { - %reduce_sum.473 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.293 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.300 = f32[]{:T(128)} add(%reduce_sum.473, %reduce_sum.293), 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":[]}} -} - -%region_158.183 (reduce_sum.459: f32[], reduce_sum.460: f32[]) -> f32[] { - %reduce_sum.459 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.460 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4162: bf16[256,512,512], param_1.5048: bf16[256,512,512]) -> (f32[], f32[]) { - %param_0.4162 = 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.4162), 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.5097 = f32[]{:T(128)} constant(0) - %reduce.672 = f32[]{:T(128)} reduce(%square.570, %constant.5097), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.5048 = 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.5048), 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.5097), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.764 = (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[] { - %reduce_sum.466 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.279 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4161: bf16[256,512,512]) -> f32[] { - %param_0.4161 = 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.4161), 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.5096 = f32[]{:T(128)} constant(0) - ROOT %reduce.673 = f32[]{:T(128)} reduce(%square.573, %constant.5096), 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[] { - %reduce_sum.935 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.631 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.632 = f32[]{:T(128)} add(%reduce_sum.935, %reduce_sum.631), 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":[]}} -} - -%region_193.218 (reduce_sum.697: f32[], reduce_sum.471: f32[]) -> f32[] { - %reduce_sum.697 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.471 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4130: f32[], param_1.5020: f32[256,1,512,512], param_2.4289: 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.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.4662.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) - %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} - %param_4.2197 = f32[]{:T(128)} parameter(4) - %div.2520.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.2519.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1361.clone.1, %div.2520.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.2519.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4830.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4252.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4830.clone.1), dimensions={}, metadata={op_name="broadcast.2362"} - %mul.4664.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %broadcast.4252.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.4829.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4251.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4829.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4663.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4251.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3453.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4664.clone.1, %mul.4663.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4289 = f32[]{:T(128)S(6)} parameter(2) - %div.2518.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4289), 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.4828.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4254.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4828.clone.1), dimensions={}, metadata={op_name="broadcast.2365"} - %mul.4666.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4254.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.1362.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\'][\'wo\']"} - %constant.4827.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4253.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4827.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4665.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4253.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3454.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_0.4130 = f32[]{:T(128)S(6)} parameter(0) - %div.2517.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4130), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2516.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3454.clone.1, %div.2517.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.2516.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4831.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4250.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4831.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3452.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.151.clone.1, %broadcast.4250.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.2518.clone.1, %add.3452.clone.1), metadata={op_name="multiply.296"} - %div.2515.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3453.clone.1, %multiply.1287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4661.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1359.clone.1, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3451.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2515.clone.1, %mul.4661.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4660.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4662.clone.1, %add.3451.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3450.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1359.clone.1, %mul.4660.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.3450.clone.1, %add.3450.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5065 = f32[]{:T(128)} constant(0) - %reduce.675 = f32[]{:T(128)} reduce(%square.577, %constant.5065), 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.3454.clone.1) - %bitcast.822.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3453.clone.1) - %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5065), 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.3450.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[] { - %reduce_sum.928 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %add.2609 = 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.21.clone.clone.clone (param_0.4928: f32[3], param_1.5502: s32[256], param_2.4341: f32[256]) -> f32[3] { + %param_0.4928 = f32[3]{0:T(128)} parameter(0) + %param_1.5502 = s32[256]{0:T(256)} parameter(1) + %reshape.4956 = s32[256]{0:T(256)} reshape(%param_1.5502), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %transpose.1166 = s32[256]{0:T(256)} transpose(%reshape.4956), dimensions={0}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} + %param_2.4341 = f32[256]{0:T(256)} parameter(2) + %reshape.4957 = f32[256]{0:T(256)} reshape(%param_2.4341) + %transpose.1167 = f32[256]{0:T(256)} transpose(%reshape.4957), dimensions={0} + ROOT %scatter-add.239 = f32[3]{0:T(128)} scatter(%param_0.4928, %transpose.1166, %transpose.1167), 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()/while/body/closed_call/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.26 (param_0.4929: f32[3], param_1.5503: s32[256], param_2.4342: f32[256]) -> f32[3] { + %param_0.4929 = f32[3]{0:T(128)} parameter(0) + %param_1.5503 = s32[256]{0:T(256)} parameter(1) + %param_2.4342 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.43 = f32[3]{0:T(128)} fusion(%param_0.4929, %param_1.5503, %param_2.4342), kind=kCustom, calls=%fused_computation.21.clone.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/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.26 (param_0.4930: f32[3], param_1.5504: s32[256], param_2.4343: f32[256]) -> f32[3] { + %param_0.4930 = f32[3]{0:T(128)} parameter(0) + %param_1.5504 = s32[256]{0:T(256)} parameter(1) + %param_2.4343 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.44.cloned.1 = f32[3]{0:T(128)} call(%param_0.4930, %param_1.5504, %param_2.4343), to_apply=%called_computation.26, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%called_computation.8 (param_0.115: f32[3], param_1.173: s32[256], param_2.118: f32[256], param_3.3125: token[]) -> f32[3] { + %param_3.3125 = 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.115 = f32[3]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %param_1.173 = 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.118 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} + %copy.2105.cloned.1.call-start = ((f32[3]{0:T(128)}), f32[3]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.115), async_execution_thread="sparsecore", calls=%async_computation.25 + %copy.2105.cloned.1.call-done = f32[3]{0:T(128)} async-done(%copy.2105.cloned.1.call-start) + %scatter_offload_custom_fusion.44.cloned.1.call-start = ((f32[3]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[3]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2105.cloned.1.call-done, %param_1.173, %param_2.118), async_execution_thread="sparsecore", calls=%async_computation.26, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} + ROOT %scatter_offload_custom_fusion.44.cloned.1.call-done = f32[3]{0:T(128)} async-done(%scatter_offload_custom_fusion.44.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%async_computation.8 (param_0.116: f32[3], param_1.174: s32[256], param_2.119: f32[256], param_3.3124: token[]) -> f32[3] { + %param_3.3124 = token[] parameter(3) + %param_0.116 = f32[3]{0:T(128)} parameter(0) + %param_1.174 = s32[256]{0:T(256)} parameter(1) + %param_2.119 = f32[256]{0:T(256)} parameter(2) + ROOT %scatter_offload_custom_fusion.26.cloned.1 = f32[3]{0:T(128)} call(%param_0.116, %param_1.174, %param_2.119, %param_3.3124), to_apply=%called_computation.8, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} +}, execution_thread="sparsecore" + +%region_145.171 (reduce_sum.321: f32[], reduce_sum.325: f32[]) -> f32[] { + %reduce_sum.321 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.325 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.326 = f32[]{:T(128)} add(%reduce_sum.321, %reduce_sum.325), 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.483 (param_0.4463: f32[3,384,128,192]) -> f32[] { + %param_0.4463 = f32[3,384,128,192]{2,3,0,1:T(8,128)} parameter(0) + %bitcast.1105 = f32[384,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4463), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.547 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.1105, %bitcast.1105), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6236 = f32[]{:T(128)} constant(0) + ROOT %reduce.698 = f32[]{:T(128)} reduce(%square.547, %constant.6236), dimensions={0,1,2,3}, to_apply=%region_145.171, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%fused_computation.484 (param_0.1460: f32[384,3,128,192]) -> bf16[3,384,128,192] { + %param_0.1460 = f32[384,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %copy.1445 = bf16[384,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1460), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\'].value"} + ROOT %bitcast.1106 = bf16[3,384,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1445), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%region_210.236 (reduce_sum.603: f32[], reduce_sum.604: f32[]) -> f32[] { + %reduce_sum.603 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.604 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.874 = f32[]{:T(128)} add(%reduce_sum.603, %reduce_sum.604), 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":[]}} +} + +%region_177.203 (reduce_sum.484: f32[], reduce_sum.485: f32[]) -> f32[] { + %reduce_sum.484 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.485 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.489 = f32[]{:T(128)} add(%reduce_sum.484, %reduce_sum.485), 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.485 (param_0.4432: f32[384,3,128,192], param_1.5188: f32[], param_2.4164: f32[], param_3.2820: f32[], param_4.2167: f32[384,3,128,192], param_5.1885: f32[], param_6.1271: f32[3,384,128,192], param_7.903: pred[], param_8.671: f32[384,3,128,192]) -> (f32[], f32[384,3,128,192], f32[384,3,128,192], f32[384,3,128,192], f32[]) { + %param_0.4432 = f32[384,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %param_3.2820 = f32[]{:T(128)S(6)} parameter(3) + %mul.4754.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2820), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.903 = pred[]{:T(512)S(6)} parameter(7) + %select_n.2205.clone.1 = pred[384,3,128,192]{2,3,1,0:T(8,128)(4,1)} broadcast(%param_7.903), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.1271 = f32[3,384,128,192]{2,3,0,1:T(8,128)} parameter(6) + %bitcast.2051.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_6.1271), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.1885 = f32[]{:T(128)} parameter(5) + %div.2622.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_5.1885), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2621.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.2051.clone.1, %div.2622.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2204.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2205.clone.1, %bitcast.2051.clone.1, %div.2621.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.5994.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4437.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.5994.clone.1), dimensions={}, metadata={op_name="broadcast.338"} + %mul.4760.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2204.clone.1, %broadcast.4437.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.671 = f32[384,3,128,192]{2,3,1,0:T(8,128)} parameter(8) + %constant.5998.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4761.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.5998.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4759.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.671, %mul.4761.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3644.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4760.clone.1, %mul.4759.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4164 = f32[]{:T(128)S(6)} parameter(2) + %div.2618.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4164), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.388.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2204.clone.1, %select_n.2204.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.5997.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4758.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.5997.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4756.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.388.clone.1, %mul.4758.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.2167 = f32[384,3,128,192]{2,3,1,0:T(8,128)} parameter(4) + %constant.5996.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4757.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.5996.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4755.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2167, %mul.4757.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3643.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4756.clone.1, %mul.4755.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5188 = f32[]{:T(128)S(6)} parameter(1) + %div.2617.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5188), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2616.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3643.clone.1, %div.2617.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.160.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} sqrt(%div.2616.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.5995.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3642.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.5995.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3641.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.160.clone.1, %add.3642.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1419.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%div.2618.clone.1, %add.3641.clone.1), metadata={op_name="multiply.288"} + %div.2615.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3644.clone.1, %multiply.1419.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4753.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4432, %broadcast.4437.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3640.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} add(%div.2615.clone.1, %mul.4753.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4752.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.4754.clone.1, %add.3640.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3639.clone.1 = f32[384,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4432, %mul.4752.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.548 = f32[384,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3639.clone.1, %add.3639.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6205 = f32[]{:T(128)} constant(0) + %reduce.699 = f32[]{:T(128)} reduce(%square.548, %constant.6205), dimensions={0,1,2,3}, to_apply=%region_210.236, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.700.clone.1 = f32[]{:T(128)} reduce(%integer_pow.388.clone.1, %constant.6205), dimensions={0,1,2,3}, to_apply=%region_177.203, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.635 = (f32[]{:T(128)}, f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.699, %add.3639.clone.1, %add.3643.clone.1, %add.3644.clone.1, %reduce.700.clone.1) +} + +%fused_computation.495 (param_0.4495: f32[32]) -> (f32[163840,32], f32[163840,32]) { + %mul.3557 = f32[163840,32]{1,0:T(8,128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_0.4495 = f32[32]{0:T(128)S(1)} parameter(0) + %mul.3650 = f32[163840,32]{1,0:T(8,128)} broadcast(%param_0.4495), dimensions={1}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.3556 = f32[163840,32]{1,0:T(8,128)} multiply(%mul.3557, %mul.3650), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.6276 = f32[]{:T(128)} constant(0) + %convert_element_type.2512 = f32[163840,32]{1,0:T(8,128)} broadcast(%constant.6276), dimensions={}, metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %exp.480 = pred[163840,32]{1,0:T(8,128)(4,1)} compare(%mul.3556, %convert_element_type.2512), direction=EQ, metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %mul.3537 = f32[163840,32]{1,0:T(8,128)} multiply(%mul.3556, %convert_element_type.2512), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %exp.473 = f32[163840,32]{1,0:T(8,128)} exponential(%mul.3537), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %constant.5671 = f32[]{:T(128)} constant(inf) + %broadcast.3966 = f32[163840,32]{1,0:T(8,128)} broadcast(%constant.5671), dimensions={}, metadata={op_name="broadcast.373"} + %exp.472 = pred[163840,32]{1,0:T(8,128)(4,1)} compare(%exp.473, %broadcast.3966), direction=EQ, metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.460 = f32[163840,32]{1,0:T(8,128)} sine(%mul.3556), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.459 = f32[163840,32]{1,0:T(8,128)} multiply(%exp.473, %exp.460), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.458 = f32[163840,32]{1,0:T(8,128)} multiply(%exp.459, %exp.473), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.457 = f32[163840,32]{1,0:T(8,128)} select(%exp.472, %exp.458, %exp.459), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.453 = f32[163840,32]{1,0:T(8,128)} select(%exp.480, %convert_element_type.2512, %exp.457), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.461.clone.1 = f32[163840,32]{1,0:T(8,128)} cosine(%mul.3556), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.452.clone.1 = f32[163840,32]{1,0:T(8,128)} multiply(%exp.473, %exp.461.clone.1), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.451.clone.1 = f32[163840,32]{1,0:T(8,128)} multiply(%exp.452.clone.1, %exp.473), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.450.clone.1 = f32[163840,32]{1,0:T(8,128)} select(%exp.472, %exp.451.clone.1, %exp.452.clone.1), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + ROOT %tuple.712 = (f32[163840,32]{1,0:T(8,128)}, f32[163840,32]{1,0:T(8,128)}) tuple(%exp.453, %exp.450.clone.1) +} + +%region_150.176 (reduce_sum.349: f32[], reduce_sum.350: f32[]) -> f32[] { + %reduce_sum.349 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.350 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.351 = f32[]{:T(128)} add(%reduce_sum.349, %reduce_sum.350), 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.4458: bf16[256,512,512], param_1.5203: s32[]) -> f32[] { + %param_0.4458 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %constant.5130 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/shard_map/jit(gmm)" stack_frame_id=0} + %param_1.5203 = s32[]{:T(128)S(6)} parameter(1) + %dynamic-slice.388 = bf16[256,512,128]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.4458, %constant.5130, %constant.5130, %param_1.5203), dynamic_slice_sizes={256,512,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294966911","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %broadcast_in_dim.1293 = f32[256,512,128]{2,1,0:T(8,128)} convert(%dynamic-slice.388), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1128 = f32[256,1,512,128]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1293), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.553 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%bitcast.1128, %bitcast.1128), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6231 = f32[]{:T(128)} constant(0) + ROOT %reduce.701 = f32[]{:T(128)} reduce(%square.553, %constant.6231), dimensions={0,1,2,3}, to_apply=%region_150.176, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_149.175 (reduce_sum.343: f32[], reduce_sum.344: f32[]) -> f32[] { + %reduce_sum.343 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.344 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.345 = f32[]{:T(128)} add(%reduce_sum.343, %reduce_sum.344), 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.4459: bf16[256,128,512]) -> f32[] { + %param_0.4459 = bf16[256,128,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %broadcast_in_dim.1297 = f32[256,128,512]{2,1,0:T(8,128)} convert(%param_0.4459), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1132 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1297), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.556 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1132, %bitcast.1132), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6232 = f32[]{:T(128)} constant(0) + ROOT %reduce.702 = f32[]{:T(128)} reduce(%square.556, %constant.6232), dimensions={0,1,2,3}, to_apply=%region_149.175, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_148.174 (reduce_sum.444: f32[], reduce_sum.445: f32[]) -> f32[] { + %reduce_sum.444 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.445 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.446 = f32[]{:T(128)} add(%reduce_sum.444, %reduce_sum.445), 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.4460: bf16[256,128,512]) -> f32[] { + %param_0.4460 = bf16[256,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %broadcast_in_dim.1301 = f32[256,128,512]{2,1,0:T(8,128)} convert(%param_0.4460), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1136 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1301), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.559 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1136, %bitcast.1136), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6233 = f32[]{:T(128)} constant(0) + ROOT %reduce.703 = f32[]{:T(128)} reduce(%square.559, %constant.6233), dimensions={0,1,2,3}, to_apply=%region_148.174, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%fused_computation.523 (param_0.1549: f32[256,1,512,128]) -> bf16[256,1,512,128] { + %param_0.1549 = f32[256,1,512,128]{3,2,1,0:T(8,128)} parameter(0) + %bitcast.1139 = f32[256,1,512,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1549), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\'].value"} + ROOT %convert_element_type.2470 = bf16[256,1,512,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.1139), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%fused_computation.528 (param_0.1561: f32[256,1,128,512]) -> bf16[256,1,128,512] { + %param_0.1561 = f32[256,1,128,512]{3,2,1,0:T(8,128)} parameter(0) + %bitcast.1148 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%param_0.1561), sharding={devices=[1,1,4,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\'].value"} + ROOT %convert_element_type.2471 = bf16[256,1,128,512]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.1148), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%fused_computation.533 (param_0.1573: f32[256,1,128,512]) -> bf16[256,1,128,512] { + %param_0.1573 = f32[256,1,128,512]{3,2,1,0:T(8,128)} parameter(0) + %bitcast.1157 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%param_0.1573), sharding={devices=[1,1,4,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\'].value"} + ROOT %convert_element_type.2472 = bf16[256,1,128,512]{3,2,0,1:T(8,128)(2,1)S(1)} convert(%bitcast.1157), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%region_216.242 (reduce_sum.624: f32[], reduce_sum.625: f32[]) -> f32[] { + %reduce_sum.624 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.625 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.626 = f32[]{:T(128)} add(%reduce_sum.928, %reduce_sum.625), 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":[]}} -} - -%region_192.217 (reduce_sum.690: f32[], reduce_sum.465: f32[]) -> f32[] { - %reduce_sum.690 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.465 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4131: f32[], param_1.5021: f32[256,1,512,512], param_2.4290: 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.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.4669.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) - %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} - %param_4.2198 = f32[]{:T(128)} parameter(4) - %div.2526.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.2525.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1365.clone.1, %div.2526.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.2525.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4835.clone.1 = f32[]{:T(128)} constant(0.1) - %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.2362"} - %mul.4671.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4257.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.4834.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4256.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4834.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4670.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3458.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4671.clone.1, %mul.4670.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4290 = f32[]{:T(128)S(6)} parameter(2) - %div.2524.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.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.4833.clone.1 = f32[]{:T(128)} constant(0.05) - %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.2365"} - %mul.4673.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4259.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.1366.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_1\']"} - %constant.4832.clone.1 = f32[]{:T(128)} constant(0.95) - %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.312"} - %mul.4672.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4258.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3459.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_0.4131 = f32[]{:T(128)S(6)} parameter(0) - %div.2523.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4131), 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(%add.3459.clone.1, %div.2523.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.2522.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4836.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4255.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4836.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3457.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.152.clone.1, %broadcast.4255.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.2524.clone.1, %add.3457.clone.1), metadata={op_name="multiply.295"} - %div.2521.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3458.clone.1, %multiply.1288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4668.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1363.clone.1, %broadcast.4257.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3456.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2521.clone.1, %mul.4668.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(%mul.4669.clone.1, %add.3456.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3455.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1363.clone.1, %mul.4667.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.3455.clone.1, %add.3455.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5066 = f32[]{:T(128)} constant(0) - %reduce.676 = f32[]{:T(128)} reduce(%square.578, %constant.5066), 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.3459.clone.1) - %bitcast.813.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3458.clone.1) - %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5066), 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.3455.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[] { - %reduce_sum.921 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.619 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.624 = f32[]{:T(128)} add(%reduce_sum.921, %reduce_sum.619), 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":[]}} -} - -%region_191.216 (reduce_sum.683: f32[], reduce_sum.463: f32[]) -> f32[] { - %reduce_sum.683 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.463 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4132: f32[], param_1.5022: f32[256,1,512,512], param_2.4291: 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.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.4676.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) - %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} - %param_4.2199 = f32[]{:T(128)} parameter(4) - %div.2532.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.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1369.clone.1, %div.2532.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.2531.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4840.clone.1 = f32[]{:T(128)} constant(0.1) - %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.2362"} - %mul.4678.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %broadcast.4262.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.4839.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4261.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4839.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.1368.clone.1, %broadcast.4261.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3463.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.4291 = f32[]{:T(128)S(6)} parameter(2) - %div.2530.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.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.4838.clone.1 = f32[]{:T(128)} constant(0.05) - %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.2365"} - %mul.4680.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5022 = 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.5022), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4837.clone.1 = f32[]{:T(128)} constant(0.95) - %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.312"} - %mul.4679.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1370.clone.1, %broadcast.4263.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3464.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.4132 = f32[]{:T(128)S(6)} parameter(0) - %div.2529.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4132), 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(%add.3464.clone.1, %div.2529.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.2528.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4841.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4260.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3462.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.153.clone.1, %broadcast.4260.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.2530.clone.1, %add.3462.clone.1), metadata={op_name="multiply.294"} - %div.2527.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3463.clone.1, %multiply.1289.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.1367.clone.1, %broadcast.4262.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3461.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2527.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.3461.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3460.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1367.clone.1, %mul.4674.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.3460.clone.1, %add.3460.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5067 = f32[]{:T(128)} constant(0) - %reduce.677 = f32[]{:T(128)} reduce(%square.579, %constant.5067), 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.3464.clone.1) - %bitcast.804.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3463.clone.1) - %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5067), 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.3460.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[] { - %reduce_sum.438 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.259 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4075: bf16[4,128,129280], param_1.4954: s32[4,128], param_2.4222: 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.4889 = 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.4888 = 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.4075 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.3151 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4075), 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.3151, %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.4887 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4888, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.4222 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.2685 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4222), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.2684 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.4887, %div.2685), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.4954 = 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.4954), 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.3150 = 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.2684, %convert_element_type.3150), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.4886 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4889, %sub.789), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.3149 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.4886), 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.935.clone.clone (param_0.4076: f32[4,128], param_1.4955: bf16[4,128,512], param_2.4224: bf16[512]) -> bf16[4,128,512] { - %param_2.4224 = 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.4224), 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.4955 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3153 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4955), 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.4076 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.4891 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4076), 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.4890 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3153, %mul.4891), 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.3152 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4890), 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.3152), 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.4165: bf16[4,128,129280], param_1.5050: s32[4,128], param_2.4316: 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.573.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.935.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.4165 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.5050 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4316 = 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.4165, %param_1.5050, %param_2.4316, %param_3.2969, %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.573.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.2653 = 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.2653, %convert_element_type.2653), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5100 = f32[]{:T(128)} constant(0) - %reduce.678 = f32[]{:T(128)} reduce(%square.581, %constant.5100), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.754 = (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[] { - %reduce_sum.564 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.387 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4149: bf16[129280,512]) -> f32[] { - %param_0.4149 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2655 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4149), 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.2655, %convert_element_type.2655), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5084 = f32[]{:T(128)} constant(0) - ROOT %reduce.679 = f32[]{:T(128)} reduce(%square.583, %constant.5084), 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[] { - %reduce_sum.1026 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.689 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.694 = f32[]{:T(128)} add(%reduce_sum.1026, %reduce_sum.689), 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":[]}} -} - -%region_206.231 (reduce_sum.788: f32[], reduce_sum.533: f32[]) -> f32[] { - %reduce_sum.788 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.533 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4117: f32[129280,512], param_1.5007: f32[], param_2.4276: 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.4117 = f32[129280,512]{1,0:T(8,128)} parameter(0) - %param_3.2932 = f32[]{:T(128)S(6)} parameter(3) - %mul.4550.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.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.3094.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.2426.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.2425.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%convert_element_type.3094.clone.1, %div.2426.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.3094.clone.1, %div.2425.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4750.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4202.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4750.clone.1), dimensions={}, metadata={op_name="broadcast.318"} - %mul.4556.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %broadcast.4202.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.4754.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4557.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4754.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4555.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4557.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3383.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4556.clone.1, %mul.4555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4276 = f32[]{:T(128)S(6)} parameter(2) - %div.2422.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4276), 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.4753.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4554.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4753.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4552.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4554.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.4752.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4553.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4752.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4551.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4553.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3382.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4552.clone.1, %mul.4551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5007 = f32[]{:T(128)S(6)} parameter(1) - %div.2421.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5007), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2420.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3382.clone.1, %div.2421.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.2420.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4751.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3381.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4751.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3380.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%sqrt.138.clone.1, %add.3381.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.2422.clone.1, %add.3380.clone.1), metadata={op_name="multiply.309"} - %div.2419.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3383.clone.1, %multiply.1274.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4549.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4117, %broadcast.4202.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3379.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2419.clone.1, %mul.4549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4548.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4550.clone.1, %add.3379.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3378.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4117, %mul.4548.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.3378.clone.1, %add.3378.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5052 = f32[]{:T(128)} constant(0) - %reduce.680 = f32[]{:T(128)} reduce(%square.584, %constant.5052), 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.5052), 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.3378.clone.1, %add.3382.clone.1, %add.3383.clone.1, %reduce.687.clone.1) -} - -%region_222.247 (reduce_sum.900: f32[], reduce_sum.605: f32[]) -> f32[] { - %reduce_sum.900 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.605 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.610 = f32[]{:T(128)} add(%reduce_sum.900, %reduce_sum.605), 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":[]}} -} - -%region_188.213 (reduce_sum.662: f32[], reduce_sum.451: f32[]) -> f32[] { - %reduce_sum.662 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.451 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4135: f32[512,129280], param_1.5025: f32[], param_2.4294: 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.4135 = f32[512,129280]{1,0:T(8,128)} parameter(0) - %param_3.2950 = f32[]{:T(128)S(6)} parameter(3) - %mul.4703.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.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.3096.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.2554.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.2553.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%convert_element_type.3096.clone.1, %div.2554.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.3096.clone.1, %div.2553.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4854.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4270.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4854.clone.1), dimensions={}, metadata={op_name="broadcast.333"} - %mul.4709.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %broadcast.4270.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.4858.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4710.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4708.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.4710.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3482.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4709.clone.1, %mul.4708.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4294 = f32[]{:T(128)S(6)} parameter(2) - %div.2550.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4294), 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.4857.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4707.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4857.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4705.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.4707.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.4856.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4706.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4856.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4704.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.4706.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3481.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4705.clone.1, %mul.4704.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5025 = f32[]{:T(128)S(6)} parameter(1) - %div.2549.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5025), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2548.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3481.clone.1, %div.2549.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.2548.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4855.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3480.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4855.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3479.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%sqrt.156.clone.1, %add.3480.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.2550.clone.1, %add.3479.clone.1), metadata={op_name="multiply.291"} - %div.2547.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3482.clone.1, %multiply.1292.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4702.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4135, %broadcast.4270.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3478.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2547.clone.1, %mul.4702.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4701.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.4703.clone.1, %add.3478.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3477.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4135, %mul.4701.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.3477.clone.1, %add.3477.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5070 = f32[]{:T(128)} constant(0) - %reduce.681 = f32[]{:T(128)} reduce(%square.585, %constant.5070), 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.5070), 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.3477.clone.1, %add.3481.clone.1, %add.3482.clone.1, %reduce.688.clone.1) -} - -%region_207.232 (reduce_sum.795: f32[], reduce_sum.535: f32[]) -> f32[] { - %reduce_sum.795 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.535 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - 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.4186: bf16[4,128,129280], param_1.5064: f32[4,128], param_2.4326: s32[4,128], param_3.2977: bf16[4,128]) -> f32[4,128] { - %param_2.4326 = 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.4326), 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.4186 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2660 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4186), 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.2660, %sub.652), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.5064 = 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.5064), 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.5124 = f32[]{:T(128)} constant(0) - %broadcast.3777 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5124), dimensions={}, metadata={op_name="broadcast.514"} - %mul.3612 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.293, %sub.639, %broadcast.3777), 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.5124), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + ROOT %reduce_sum.916 = f32[]{:T(128)} add(%reduce_sum.624, %reduce_sum.625), 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":[]}} +} + +%region_182.208 (reduce_sum.505: f32[], reduce_sum.506: f32[]) -> f32[] { + %reduce_sum.505 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.506 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.678 = f32[]{:T(128)} add(%reduce_sum.505, %reduce_sum.506), 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.536 (param_0.4427: f32[], param_1.5183: f32[256,1,512,128], param_2.4159: f32[], param_3.2815: f32[256,1,512,128], param_4.2162: f32[], param_5.1880: bf16[256,512,512], param_6.1266: s32[], param_7.898: pred[], param_8.667: f32[], param_9.587: f32[256,1,512,128]) -> (f32[], f32[256,1,512,128], f32[256,1,512,128], f32[256,1,512,128], f32[]) { + %param_9.587 = f32[256,1,512,128]{3,2,1,0:T(8,128)} parameter(9) + %bitcast.2036.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} bitcast(%param_9.587), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\'].value"} + %param_8.667 = f32[]{:T(128)S(6)} parameter(8) + %mul.4713.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%param_8.667), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.898 = pred[]{:T(512)S(6)} parameter(7) + %select_n.2189.clone.1 = pred[256,1,512,128]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_7.898), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_5.1880 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) + %constant.5966.clone.1 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/shard_map/jit(gmm)" stack_frame_id=0} + %param_6.1266 = s32[]{:T(128)S(6)} parameter(6) + %dynamic-slice.461.clone.1 = bf16[256,512,128]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_5.1880, %constant.5966.clone.1, %constant.5966.clone.1, %param_6.1266), dynamic_slice_sizes={256,512,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294966911","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %broadcast_in_dim.1511.clone.1 = f32[256,512,128]{2,1,0:T(8,128)} convert(%dynamic-slice.461.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.2038.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1511.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_4.2162 = f32[]{:T(128)} parameter(4) + %div.2586.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%param_4.2162), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2585.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} divide(%bitcast.2038.clone.1, %div.2586.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2188.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} select(%select_n.2189.clone.1, %bitcast.2038.clone.1, %div.2585.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.5970.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4419.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%constant.5970.clone.1), dimensions={}, metadata={op_name="broadcast.2453"} + %mul.4715.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%select_n.2188.clone.1, %broadcast.4419.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2815 = f32[256,1,512,128]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.2037.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} bitcast(%param_3.2815), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'optimizer\'][\'opt_state\'][0][\'mu\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\'].value"} + %constant.5969.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4418.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%constant.5969.clone.1), dimensions={}, metadata={op_name="broadcast.333"} + %mul.4714.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%bitcast.2037.clone.1, %broadcast.4418.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3615.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} add(%mul.4715.clone.1, %mul.4714.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4159 = f32[]{:T(128)S(6)} parameter(2) + %div.2584.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%param_2.4159), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.383.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%select_n.2188.clone.1, %select_n.2188.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.5968.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4421.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%constant.5968.clone.1), dimensions={}, metadata={op_name="broadcast.2456"} + %mul.4717.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%integer_pow.383.clone.1, %broadcast.4421.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5183 = f32[256,1,512,128]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.2039.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} bitcast(%param_1.5183), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'optimizer\'][\'opt_state\'][0][\'nu\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\'].value"} + %constant.5967.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4420.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%constant.5967.clone.1), dimensions={}, metadata={op_name="broadcast.316"} + %mul.4716.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%bitcast.2039.clone.1, %broadcast.4420.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3616.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} add(%mul.4717.clone.1, %mul.4716.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4427 = f32[]{:T(128)S(6)} parameter(0) + %div.2583.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%param_0.4427), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2582.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} divide(%add.3616.clone.1, %div.2583.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.155.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} sqrt(%div.2582.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.5971.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4417.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} broadcast(%constant.5971.clone.1), dimensions={}, metadata={op_name="broadcast.309"} + %add.3614.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} add(%sqrt.155.clone.1, %broadcast.4417.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1414.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%div.2584.clone.1, %add.3614.clone.1), metadata={op_name="multiply.294"} + %div.2581.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} divide(%add.3615.clone.1, %multiply.1414.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4712.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%bitcast.2036.clone.1, %broadcast.4419.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3613.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} add(%div.2581.clone.1, %mul.4712.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4711.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%mul.4713.clone.1, %add.3613.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3612.clone.1 = f32[256,1,512,128]{3,2,0,1:T(8,128)} add(%bitcast.2036.clone.1, %mul.4711.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.560 = f32[256,1,512,128]{3,2,0,1:T(8,128)} multiply(%add.3612.clone.1, %add.3612.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6200 = f32[]{:T(128)} constant(0) + %reduce.704 = f32[]{:T(128)} reduce(%square.560, %constant.6200), dimensions={0,1,2,3}, to_apply=%region_216.242, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.1295.clone.1 = f32[256,1,512,128]{3,2,1,0:T(8,128)} bitcast(%add.3616.clone.1) + %bitcast.1286.clone.1 = f32[256,1,512,128]{3,2,1,0:T(8,128)} bitcast(%add.3615.clone.1) + %reduce.713.clone.1 = f32[]{:T(128)} reduce(%integer_pow.383.clone.1, %constant.6200), dimensions={0,1,2,3}, to_apply=%region_182.208, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.634 = (f32[]{:T(128)}, f32[256,1,512,128]{3,2,0,1:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.704, %add.3612.clone.1, %bitcast.1295.clone.1, %bitcast.1286.clone.1, %reduce.713.clone.1) +} + +%region_215.241 (reduce_sum.622: f32[], reduce_sum.623: f32[]) -> f32[] { + %reduce_sum.622 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.623 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.909 = f32[]{:T(128)} add(%reduce_sum.622, %reduce_sum.623), 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":[]}} +} + +%region_181.207 (reduce_sum.503: f32[], reduce_sum.504: f32[]) -> f32[] { + %reduce_sum.503 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.504 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.671 = f32[]{:T(128)} add(%reduce_sum.503, %reduce_sum.504), 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":[]}} +} + +%region_180.206 (reduce_sum.498: f32[], reduce_sum.499: f32[]) -> f32[] { + %reduce_sum.498 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.499 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.664 = f32[]{:T(128)} add(%reduce_sum.498, %reduce_sum.499), 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":[]}} +} + +%region_214.240 (reduce_sum.617: f32[], reduce_sum.618: f32[]) -> f32[] { + %reduce_sum.617 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.618 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.902 = f32[]{:T(128)} add(%reduce_sum.617, %reduce_sum.618), 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.537 (param_0.4428: f32[], param_1.5184: f32[256,1,128,512], param_2.4160: f32[], param_3.2816: f32[256,1,128,512], param_4.2163: f32[], param_5.1881: bf16[256,128,512], param_6.1267: pred[], param_7.899: f32[], param_8.668: f32[256,1,128,512], param_9.588: bf16[256,128,512], param_10.529: f32[256,1,128,512], param_11.477: f32[256,1,128,512], param_12.349: f32[256,1,128,512]) -> (f32[], f32[256,1,128,512], f32[256,1,128,512], f32[256,1,128,512], f32[], /*index=5*/f32[], f32[256,1,128,512], f32[256,1,128,512], f32[], f32[256,1,128,512]) { + %param_8.668 = f32[256,1,128,512]{3,2,1,0:T(8,128)} parameter(8) + %bitcast.2040.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%param_8.668), sharding={devices=[1,1,4,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\'].value"} + %param_7.899 = f32[]{:T(128)S(6)} parameter(7) + %mul.4720.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%param_7.899), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_6.1267 = pred[]{:T(512)S(6)} parameter(6) + %select_n.2191.clone.1 = pred[256,1,128,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1267), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_5.1881 = bf16[256,128,512]{2,1,0:T(8,128)(2,1)} parameter(5) + %broadcast_in_dim.1512.clone.1 = f32[256,128,512]{2,1,0:T(8,128)} convert(%param_5.1881), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.2042.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1512.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_4.2163 = f32[]{:T(128)} parameter(4) + %div.2592.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2163), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2591.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} divide(%bitcast.2042.clone.1, %div.2592.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2190.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} select(%select_n.2191.clone.1, %bitcast.2042.clone.1, %div.2591.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.5975.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4424.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%constant.5975.clone.1), dimensions={}, metadata={op_name="broadcast.2453"} + %mul.4722.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%select_n.2190.clone.1, %broadcast.4424.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_3.2816 = f32[256,1,128,512]{3,2,1,0:T(8,128)} parameter(3) + %bitcast.2041.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2816), sharding={devices=[1,1,4,1]<=[4]}, metadata={op_name="state[\'optimizer\'][\'opt_state\'][0][\'mu\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\'].value"} + %constant.5974.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.4423.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%constant.5974.clone.1), dimensions={}, metadata={op_name="broadcast.333"} + %mul.4721.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%bitcast.2041.clone.1, %broadcast.4423.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3620.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%mul.4722.clone.1, %mul.4721.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4160 = f32[]{:T(128)S(6)} parameter(2) + %div.2590.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4160), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.384.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%select_n.2190.clone.1, %select_n.2190.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.5973.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.4426.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%constant.5973.clone.1), dimensions={}, metadata={op_name="broadcast.2456"} + %mul.4724.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.384.clone.1, %broadcast.4426.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_1.5184 = f32[256,1,128,512]{3,2,1,0:T(8,128)} parameter(1) + %bitcast.2043.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5184), sharding={devices=[1,1,4,1]<=[4]}, metadata={op_name="state[\'optimizer\'][\'opt_state\'][0][\'nu\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\'].value"} + %constant.5972.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.4425.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%constant.5972.clone.1), dimensions={}, metadata={op_name="broadcast.316"} + %mul.4723.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%bitcast.2043.clone.1, %broadcast.4425.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3621.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%mul.4724.clone.1, %mul.4723.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4428 = f32[]{:T(128)S(6)} parameter(0) + %div.2589.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4428), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2588.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} divide(%add.3621.clone.1, %div.2589.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.156.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} sqrt(%div.2588.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.5976.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.4422.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} broadcast(%constant.5976.clone.1), dimensions={}, metadata={op_name="broadcast.309"} + %add.3619.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%sqrt.156.clone.1, %broadcast.4422.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1415.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%div.2590.clone.1, %add.3619.clone.1), metadata={op_name="multiply.293"} + %div.2587.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} divide(%add.3620.clone.1, %multiply.1415.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4719.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%bitcast.2040.clone.1, %broadcast.4424.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3618.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%div.2587.clone.1, %mul.4719.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4718.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%mul.4720.clone.1, %add.3618.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3617.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%bitcast.2040.clone.1, %mul.4718.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.561 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%add.3617.clone.1, %add.3617.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6201 = f32[]{:T(128)} constant(0) + %reduce.705 = f32[]{:T(128)} reduce(%square.561, %constant.6201), dimensions={0,1,2,3}, to_apply=%region_215.241, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %bitcast.1277.clone.1 = f32[256,1,128,512]{3,2,1,0:T(8,128)} bitcast(%add.3621.clone.1) + %bitcast.1259.clone.1 = f32[256,1,128,512]{3,2,1,0:T(8,128)} bitcast(%add.3620.clone.1) + %reduce.714.clone.1 = f32[]{:T(128)} reduce(%integer_pow.384.clone.1, %constant.6201), dimensions={0,1,2,3}, to_apply=%region_181.207, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_9.588 = bf16[256,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(9) + %broadcast_in_dim.1303.clone.1.clone.1 = f32[256,128,512]{2,1,0:T(8,128)} convert(%param_9.588), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1186.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1303.clone.1.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %div.1767.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} divide(%bitcast.1186.clone.1.clone.1, %div.2592.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.1558.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} select(%select_n.2191.clone.1, %bitcast.1186.clone.1.clone.1, %div.1767.clone.1.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %integer_pow.295.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%select_n.1558.clone.1.clone.1, %select_n.1558.clone.1.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %reduce.715.clone.1.clone.1 = f32[]{:T(128)} reduce(%integer_pow.295.clone.1.clone.1, %constant.6201), dimensions={0,1,2,3}, to_apply=%region_180.206, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %mul.3709.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%select_n.1558.clone.1.clone.1, %broadcast.4424.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_10.529 = f32[256,1,128,512]{3,2,1,0:T(8,128)} parameter(10) + %bitcast.1253.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%param_10.529), sharding={devices=[1,1,4,1]<=[4]}, metadata={op_name="state[\'optimizer\'][\'opt_state\'][0][\'mu\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\'].value"} + %mul.3708.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1253.clone.1.clone.1, %broadcast.4423.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3083.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%mul.3709.clone.1.clone.1, %mul.3708.clone.1.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %bitcast.1250.clone.1.clone.1 = f32[256,1,128,512]{3,2,1,0:T(8,128)} bitcast(%add.3083.clone.1.clone.1) + %mul.3725.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.295.clone.1.clone.1, %broadcast.4426.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_11.477 = f32[256,1,128,512]{3,2,1,0:T(8,128)} parameter(11) + %bitcast.1271.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%param_11.477), sharding={devices=[1,1,4,1]<=[4]}, metadata={op_name="state[\'optimizer\'][\'opt_state\'][0][\'nu\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\'].value"} + %mul.3724.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1271.clone.1.clone.1, %broadcast.4425.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3091.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%mul.3725.clone.1.clone.1, %mul.3724.clone.1.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %bitcast.1268.clone.1.clone.1 = f32[256,1,128,512]{3,2,1,0:T(8,128)} bitcast(%add.3091.clone.1.clone.1) + %param_12.349 = f32[256,1,128,512]{3,2,1,0:T(8,128)} parameter(12) + %bitcast.2044.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} bitcast(%param_12.349), sharding={devices=[1,1,4,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\'].value"} + %div.2594.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} divide(%add.3091.clone.1.clone.1, %div.2589.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.157.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} sqrt(%div.2594.clone.1.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %add.3624.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%sqrt.157.clone.1.clone.1, %broadcast.4422.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1416.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%div.2590.clone.1, %add.3624.clone.1.clone.1), metadata={op_name="multiply.292"} + %div.2593.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} divide(%add.3083.clone.1.clone.1, %multiply.1416.clone.1.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4726.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%bitcast.2044.clone.1.clone.1, %broadcast.4424.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3623.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%div.2593.clone.1.clone.1, %mul.4726.clone.1.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4725.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%mul.4720.clone.1, %add.3623.clone.1.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3622.clone.1.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} add(%bitcast.2044.clone.1.clone.1, %mul.4725.clone.1.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.562.clone.1 = f32[256,1,128,512]{3,2,0,1:T(8,128)} multiply(%add.3622.clone.1.clone.1, %add.3622.clone.1.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.706.clone.1 = f32[]{:T(128)} reduce(%square.562.clone.1, %constant.6201), dimensions={0,1,2,3}, to_apply=%region_214.240, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.730 = (f32[]{:T(128)}, f32[256,1,128,512]{3,2,0,1:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}, /*index=5*/f32[]{:T(128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}, f32[256,1,128,512]{3,2,0,1:T(8,128)}) tuple(%reduce.705, %add.3617.clone.1, %bitcast.1277.clone.1, %bitcast.1259.clone.1, %reduce.714.clone.1, /*index=5*/%reduce.715.clone.1.clone.1, %bitcast.1250.clone.1.clone.1, %bitcast.1268.clone.1.clone.1, %reduce.706.clone.1, %add.3622.clone.1.clone.1) +} + +%region_146.172 (reduce_sum.327: f32[], reduce_sum.328: f32[]) -> f32[] { + %reduce_sum.327 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.328 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.332 = f32[]{:T(128)} add(%reduce_sum.327, %reduce_sum.328), 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.539 (param_0.4462: bf16[128,129280]) -> f32[] { + %param_0.4462 = bf16[128,129280]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2474 = f32[128,129280]{1,0:T(8,128)} convert(%param_0.4462), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %square.564 = f32[128,129280]{1,0:T(8,128)} multiply(%convert_element_type.2474, %convert_element_type.2474), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6235 = f32[]{:T(128)} constant(0) + ROOT %reduce.707 = f32[]{:T(128)} reduce(%square.564, %constant.6235), dimensions={0,1}, to_apply=%region_146.172, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_164.190 (reduce_sum.415: f32[], reduce_sum.419: f32[]) -> f32[] { + %reduce_sum.415 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.419 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.420 = f32[]{:T(128)} add(%reduce_sum.415, %reduce_sum.419), 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.4448: bf16[129280,128]) -> f32[] { + %param_0.4448 = bf16[129280,128]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2476 = f32[129280,128]{1,0:T(8,128)} convert(%param_0.4448), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %square.566 = f32[129280,128]{1,0:T(8,128)} multiply(%convert_element_type.2476, %convert_element_type.2476), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6221 = f32[]{:T(128)} constant(0) + ROOT %reduce.708 = f32[]{:T(128)} reduce(%square.566, %constant.6221), dimensions={0,1}, to_apply=%region_164.190, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_37.47 (reduce_sum.76: f32[], reduce_sum.80: f32[]) -> f32[] { @@ -1850,151 +1551,450 @@ 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.4187: bf16[4,128,129280], param_1.5065: bf16[4,128]) -> f32[4,128] { - %param_0.4187 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2666 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4187), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.5065 = 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.5065), 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.2666, %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.5125 = f32[]{:T(128)} constant(0) - ROOT %reduce.683 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.448, %constant.5125), 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[] { - %reduce_sum.417 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.244 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4168: f32[3,512,128,256]) -> f32[] { - %param_0.4168 = 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.4168), 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.5103 = f32[]{:T(128)} constant(0) - ROOT %reduce.689 = f32[]{:T(128)} reduce(%square.588, %constant.5103), 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.1601: f32[512,3,128,256]) -> bf16[3,512,128,256] { - %param_0.1601 = 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.1601), 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[] { - %reduce_sum.879 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.591 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.596 = f32[]{:T(128)} add(%reduce_sum.879, %reduce_sum.591), 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":[]}} -} - -%region_185.210 (reduce_sum.641: f32[], reduce_sum.437: f32[]) -> f32[] { - %reduce_sum.641 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.437 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4138: f32[512,3,128,256], param_1.5028: f32[], param_2.4297: 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.4138 = 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.4733.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.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} - %param_5.2008 = f32[]{:T(128)} parameter(5) - %div.2578.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.2577.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.1378.clone.1, %div.2578.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.2577.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4872.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4276.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4872.clone.1), dimensions={}, metadata={op_name="broadcast.336"} - %mul.4739.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %broadcast.4276.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.4876.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4740.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)/mul" stack_frame_id=0} - %mul.4738.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.4740.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3500.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4739.clone.1, %mul.4738.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4297 = f32[]{:T(128)S(6)} parameter(2) - %div.2574.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4297), 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.4875.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4737.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4875.clone.1), 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)} multiply(%integer_pow.401.clone.1, %mul.4737.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.4874.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4736.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4874.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4734.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.4736.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3499.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4735.clone.1, %mul.4734.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5028 = f32[]{:T(128)S(6)} parameter(1) - %div.2573.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5028), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2572.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3499.clone.1, %div.2573.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.2572.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4873.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3498.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4873.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3497.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3498.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.2574.clone.1, %add.3497.clone.1), metadata={op_name="multiply.288"} - %div.2571.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3500.clone.1, %multiply.1295.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4732.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4138, %broadcast.4276.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3496.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2571.clone.1, %mul.4732.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4731.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.4733.clone.1, %add.3496.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3495.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4138, %mul.4731.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.3495.clone.1, %add.3495.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5073 = f32[]{:T(128)} constant(0) - %reduce.690 = f32[]{:T(128)} reduce(%square.589, %constant.5073), 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.5073), 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.3495.clone.1, %add.3499.clone.1, %add.3500.clone.1, %reduce.691.clone.1) -} - -%region_172.197 (reduce_sum.557: f32[], reduce_sum.381: f32[]) -> f32[] { - %reduce_sum.557 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.381 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4102: f32[4,128], param_1.4999: bf16[4,128,1536], param_2.4258: bf16[1536]) -> bf16[4,128,1536,1] { - %param_2.4258 = 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.4258), 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.4999 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3175 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4999), 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.4102 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.4937 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4102), 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.4936 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3175, %mul.4937), 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.3174 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.4936), 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.3174), 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.1488 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} bitcast(%bitcast_input.12) -} - -%fused_computation.551 (param_0.4150: bf16[4,128,128,192], param_1.5039: f32[4,128], param_2.4308: bf16[4,128,1536], param_3.2964: bf16[1536]) -> (f32[], bf16[1536,128,192,1]) { - %param_1.5039 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4308 = 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.459.clone.1 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.5039, %param_2.4308, %param_3.2964), 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.4150 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - %fusion.746 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.4150), 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.459.clone.1, %fusion.746), 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.5085 = f32[]{:T(128)} constant(0) - %reduce.692 = f32[]{:T(128)} reduce(%square.592, %constant.5085), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.763 = (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[] { - %reduce_sum.1019 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.687 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.688 = f32[]{:T(128)} add(%reduce_sum.1019, %reduce_sum.687), 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":[]}} -} - -%region_205.230 (reduce_sum.781: f32[], reduce_sum.527: f32[]) -> f32[] { - %reduce_sum.781 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.527 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - 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.4118: f32[], param_1.5008: f32[], param_2.4277: 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.545 (param_0.4484: bf16[128,129280], param_1.5216: bf16[128]) -> f32[128] { + %param_0.4484 = bf16[128,129280]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2485 = f32[128,129280]{1,0:T(8,128)} convert(%param_0.4484), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.1172 = f32[1,128,129280]{2,1,0:T(8,128)} bitcast(%convert_element_type.2485), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.5216 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %sub.691 = f32[1,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5216), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.662 = f32[1,128,129280]{2,1,0:T(8,128)} subtract(%bitcast.1172, %sub.691), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.485 = f32[1,128,129280]{2,1,0:T(8,128)} exponential(%sub.662), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %constant.6260 = f32[]{:T(128)} constant(0) + ROOT %reduce.709 = f32[128]{0:T(128)S(1)} reduce(%exp.485, %constant.6260), dimensions={0,2}, to_apply=%region_37.47, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%region_196.222 (reduce_sum.554: f32[], reduce_sum.555: f32[]) -> f32[] { + %reduce_sum.554 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.555 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.776 = f32[]{:T(128)} add(%reduce_sum.554, %reduce_sum.555), 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.551 (param_0.4483: bf16[128,129280], param_1.5214: f32[128], param_2.4183: bf16[128], param_3.2836: s32[128]) -> f32[128] { + %param_3.2836 = s32[128]{0:T(128)S(1)} parameter(3) + %eq.311 = s32[1,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2836), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.289 = s32[1,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.306 = pred[1,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.311, %iota.289), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %param_0.4483 = bf16[128,129280]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2498 = f32[128,129280]{1,0:T(8,128)} convert(%param_0.4483), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.1184 = f32[1,128,129280]{2,1,0:T(8,128)} bitcast(%convert_element_type.2498), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_2.4183 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %sub.690 = f32[1,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4183), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.683 = f32[1,128,129280]{2,1,0:T(8,128)} subtract(%bitcast.1184, %sub.690), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.5214 = f32[128]{0:T(128)S(1)} parameter(1) + %sub.688 = f32[1,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5214), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.682 = f32[1,128,129280]{2,1,0:T(8,128)} subtract(%sub.683, %sub.688), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %constant.6258 = f32[]{:T(128)} constant(0) + %broadcast.3986 = f32[1,128,129280]{2,1,0:T(8,128)} broadcast(%constant.6258), dimensions={}, metadata={op_name="broadcast.563"} + %mul.3626 = f32[1,128,129280]{2,1,0:T(8,128)} select(%eq.306, %sub.682, %broadcast.3986), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.710 = f32[128]{0:T(128)S(1)} reduce(%mul.3626, %constant.6258), dimensions={0,2}, to_apply=%region_196.222, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%region_229.255 (reduce_sum.674: f32[], reduce_sum.679: f32[]) -> f32[] { + %reduce_sum.674 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.679 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.1007 = f32[]{:T(128)} add(%reduce_sum.674, %reduce_sum.679), 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":[]}} +} + +%region_195.221 (reduce_sum.552: f32[], reduce_sum.553: f32[]) -> f32[] { + %reduce_sum.552 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.553 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.769 = f32[]{:T(128)} add(%reduce_sum.552, %reduce_sum.553), 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.552 (param_0.4416: f32[129280,128], param_1.5172: f32[], param_2.4148: f32[], param_3.2804: f32[], param_4.2151: f32[129280,128], param_5.1869: f32[], param_6.1255: bf16[129280,128], param_7.887: pred[], param_8.656: f32[129280,128]) -> (f32[], f32[129280,128], f32[129280,128], f32[129280,128], f32[]) { + %param_0.4416 = f32[129280,128]{1,0:T(8,128)} parameter(0) + %param_3.2804 = f32[]{:T(128)S(6)} parameter(3) + %mul.4622.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%param_3.2804), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.887 = pred[]{:T(512)S(6)} parameter(7) + %select_n.2153.clone.1 = pred[129280,128]{1,0:T(8,128)(4,1)} broadcast(%param_7.887), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.1255 = bf16[129280,128]{1,0:T(8,128)(2,1)S(1)} parameter(6) + %convert_element_type.2854.clone.1 = f32[129280,128]{1,0:T(8,128)} convert(%param_6.1255), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_5.1869 = f32[]{:T(128)} parameter(5) + %div.2510.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%param_5.1869), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2509.clone.1 = f32[129280,128]{1,0:T(8,128)} divide(%convert_element_type.2854.clone.1, %div.2510.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2152.clone.1 = f32[129280,128]{1,0:T(8,128)} select(%select_n.2153.clone.1, %convert_element_type.2854.clone.1, %div.2509.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.5899.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4387.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%constant.5899.clone.1), dimensions={}, metadata={op_name="broadcast.322"} + %mul.4628.clone.1 = f32[129280,128]{1,0:T(8,128)} multiply(%select_n.2152.clone.1, %broadcast.4387.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.656 = f32[129280,128]{1,0:T(8,128)} parameter(8) + %constant.5903.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4629.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%constant.5903.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4627.clone.1 = f32[129280,128]{1,0:T(8,128)} multiply(%param_8.656, %mul.4629.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3560.clone.1 = f32[129280,128]{1,0:T(8,128)} add(%mul.4628.clone.1, %mul.4627.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4148 = f32[]{:T(128)S(6)} parameter(2) + %div.2506.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%param_2.4148), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.373.clone.1 = f32[129280,128]{1,0:T(8,128)} multiply(%select_n.2152.clone.1, %select_n.2152.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.5902.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4626.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%constant.5902.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4624.clone.1 = f32[129280,128]{1,0:T(8,128)} multiply(%integer_pow.373.clone.1, %mul.4626.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.2151 = f32[129280,128]{1,0:T(8,128)} parameter(4) + %constant.5901.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4625.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%constant.5901.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4623.clone.1 = f32[129280,128]{1,0:T(8,128)} multiply(%param_4.2151, %mul.4625.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3559.clone.1 = f32[129280,128]{1,0:T(8,128)} add(%mul.4624.clone.1, %mul.4623.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5172 = f32[]{:T(128)S(6)} parameter(1) + %div.2505.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%param_1.5172), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2504.clone.1 = f32[129280,128]{1,0:T(8,128)} divide(%add.3559.clone.1, %div.2505.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.145.clone.1 = f32[129280,128]{1,0:T(8,128)} sqrt(%div.2504.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.5900.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3558.clone.1 = f32[129280,128]{1,0:T(8,128)} broadcast(%constant.5900.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3557.clone.1 = f32[129280,128]{1,0:T(8,128)} add(%sqrt.145.clone.1, %add.3558.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1404.clone.1 = f32[129280,128]{1,0:T(8,128)} multiply(%div.2506.clone.1, %add.3557.clone.1), metadata={op_name="multiply.307"} + %div.2503.clone.1 = f32[129280,128]{1,0:T(8,128)} divide(%add.3560.clone.1, %multiply.1404.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4621.clone.1 = f32[129280,128]{1,0:T(8,128)} multiply(%param_0.4416, %broadcast.4387.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3556.clone.1 = f32[129280,128]{1,0:T(8,128)} add(%div.2503.clone.1, %mul.4621.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4620.clone.1 = f32[129280,128]{1,0:T(8,128)} multiply(%mul.4622.clone.1, %add.3556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3555.clone.1 = f32[129280,128]{1,0:T(8,128)} add(%param_0.4416, %mul.4620.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.567 = f32[129280,128]{1,0:T(8,128)} multiply(%add.3555.clone.1, %add.3555.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6189 = f32[]{:T(128)} constant(0) + %reduce.711 = f32[]{:T(128)} reduce(%square.567, %constant.6189), dimensions={0,1}, to_apply=%region_229.255, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.716.clone.1 = f32[]{:T(128)} reduce(%integer_pow.373.clone.1, %constant.6189), dimensions={0,1}, to_apply=%region_195.221, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.643 = (f32[]{:T(128)}, f32[129280,128]{1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.711, %add.3555.clone.1, %add.3559.clone.1, %add.3560.clone.1, %reduce.716.clone.1) +} + +%region_211.237 (reduce_sum.608: f32[], reduce_sum.609: f32[]) -> f32[] { + %reduce_sum.608 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.609 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.881 = f32[]{:T(128)} add(%reduce_sum.608, %reduce_sum.609), 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":[]}} +} + +%region_178.204 (reduce_sum.490: f32[], reduce_sum.491: f32[]) -> f32[] { + %reduce_sum.490 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.491 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.492 = f32[]{:T(128)} add(%reduce_sum.490, %reduce_sum.491), 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.553 (param_0.4431: f32[128,129280], param_1.5187: f32[], param_2.4163: f32[], param_3.2819: f32[], param_4.2166: f32[128,129280], param_5.1884: f32[], param_6.1270: bf16[128,129280], param_7.902: pred[], param_8.670: f32[128,129280]) -> (f32[], f32[128,129280], f32[128,129280], f32[128,129280], f32[]) { + %param_0.4431 = f32[128,129280]{1,0:T(8,128)} parameter(0) + %param_3.2819 = f32[]{:T(128)S(6)} parameter(3) + %mul.4744.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%param_3.2819), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.902 = pred[]{:T(512)S(6)} parameter(7) + %select_n.2201.clone.1 = pred[128,129280]{1,0:T(8,128)(4,1)} broadcast(%param_7.902), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.1270 = bf16[128,129280]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.2856.clone.1 = f32[128,129280]{1,0:T(8,128)} convert(%param_6.1270), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_5.1884 = f32[]{:T(128)} parameter(5) + %div.2614.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%param_5.1884), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2613.clone.1 = f32[128,129280]{1,0:T(8,128)} divide(%convert_element_type.2856.clone.1, %div.2614.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2200.clone.1 = f32[128,129280]{1,0:T(8,128)} select(%select_n.2201.clone.1, %convert_element_type.2856.clone.1, %div.2613.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.5988.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4435.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%constant.5988.clone.1), dimensions={}, metadata={op_name="broadcast.337"} + %mul.4750.clone.1 = f32[128,129280]{1,0:T(8,128)} multiply(%select_n.2200.clone.1, %broadcast.4435.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.670 = f32[128,129280]{1,0:T(8,128)} parameter(8) + %constant.5992.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4751.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%constant.5992.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4749.clone.1 = f32[128,129280]{1,0:T(8,128)} multiply(%param_8.670, %mul.4751.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3638.clone.1 = f32[128,129280]{1,0:T(8,128)} add(%mul.4750.clone.1, %mul.4749.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4163 = f32[]{:T(128)S(6)} parameter(2) + %div.2610.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%param_2.4163), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.387.clone.1 = f32[128,129280]{1,0:T(8,128)} multiply(%select_n.2200.clone.1, %select_n.2200.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.5991.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4748.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%constant.5991.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4746.clone.1 = f32[128,129280]{1,0:T(8,128)} multiply(%integer_pow.387.clone.1, %mul.4748.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.2166 = f32[128,129280]{1,0:T(8,128)} parameter(4) + %constant.5990.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4747.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%constant.5990.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4745.clone.1 = f32[128,129280]{1,0:T(8,128)} multiply(%param_4.2166, %mul.4747.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3637.clone.1 = f32[128,129280]{1,0:T(8,128)} add(%mul.4746.clone.1, %mul.4745.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5187 = f32[]{:T(128)S(6)} parameter(1) + %div.2609.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%param_1.5187), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2608.clone.1 = f32[128,129280]{1,0:T(8,128)} divide(%add.3637.clone.1, %div.2609.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.159.clone.1 = f32[128,129280]{1,0:T(8,128)} sqrt(%div.2608.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.5989.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3636.clone.1 = f32[128,129280]{1,0:T(8,128)} broadcast(%constant.5989.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3635.clone.1 = f32[128,129280]{1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3636.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1418.clone.1 = f32[128,129280]{1,0:T(8,128)} multiply(%div.2610.clone.1, %add.3635.clone.1), metadata={op_name="multiply.289"} + %div.2607.clone.1 = f32[128,129280]{1,0:T(8,128)} divide(%add.3638.clone.1, %multiply.1418.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4743.clone.1 = f32[128,129280]{1,0:T(8,128)} multiply(%param_0.4431, %broadcast.4435.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3634.clone.1 = f32[128,129280]{1,0:T(8,128)} add(%div.2607.clone.1, %mul.4743.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4742.clone.1 = f32[128,129280]{1,0:T(8,128)} multiply(%mul.4744.clone.1, %add.3634.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3633.clone.1 = f32[128,129280]{1,0:T(8,128)} add(%param_0.4431, %mul.4742.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.568 = f32[128,129280]{1,0:T(8,128)} multiply(%add.3633.clone.1, %add.3633.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6204 = f32[]{:T(128)} constant(0) + %reduce.712 = f32[]{:T(128)} reduce(%square.568, %constant.6204), dimensions={0,1}, to_apply=%region_211.237, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.717.clone.1 = f32[]{:T(128)} reduce(%integer_pow.387.clone.1, %constant.6204), dimensions={0,1}, to_apply=%region_178.204, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.644 = (f32[]{:T(128)}, f32[128,129280]{1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.712, %add.3633.clone.1, %add.3637.clone.1, %add.3638.clone.1, %reduce.717.clone.1) +} + +%region_143.169 (reduce_sum.282: f32[], reduce_sum.289: f32[]) -> f32[] { + %reduce_sum.282 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.289 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.296 = f32[]{:T(128)} add(%reduce_sum.282, %reduce_sum.289), 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.568 (param_0.4465: f32[3,128,128,256]) -> f32[] { + %param_0.4465 = f32[3,128,128,256]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.1206 = f32[128,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4465), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.571 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.1206, %bitcast.1206), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6238 = f32[]{:T(128)} constant(0) + ROOT %reduce.718 = f32[]{:T(128)} reduce(%square.571, %constant.6238), dimensions={0,1,2,3}, to_apply=%region_143.169, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%fused_computation.569 (param_0.1678: f32[128,3,128,256]) -> bf16[3,128,128,256] { + %param_0.1678 = f32[128,3,128,256]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.1446 = bf16[128,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1678), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\'].value"} + ROOT %bitcast.1207 = bf16[3,128,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1446), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%region_208.234 (reduce_sum.596: f32[], reduce_sum.597: f32[]) -> f32[] { + %reduce_sum.596 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.597 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.860 = f32[]{:T(128)} add(%reduce_sum.596, %reduce_sum.597), 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":[]}} +} + +%region_175.201 (reduce_sum.475: f32[], reduce_sum.476: f32[]) -> f32[] { + %reduce_sum.475 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.476 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.477 = f32[]{:T(128)} add(%reduce_sum.475, %reduce_sum.476), 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.570 (param_0.4434: f32[128,3,128,256], param_1.5190: f32[], param_2.4166: f32[], param_3.2822: f32[], param_4.2169: f32[128,3,128,256], param_5.1887: f32[], param_6.1273: f32[3,128,128,256], param_7.905: pred[], param_8.673: f32[128,3,128,256]) -> (f32[], f32[128,3,128,256], f32[128,3,128,256], f32[128,3,128,256], f32[]) { + %param_0.4434 = f32[128,3,128,256]{3,2,1,0:T(8,128)} parameter(0) + %param_3.2822 = f32[]{:T(128)S(6)} parameter(3) + %mul.4774.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2822), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.905 = pred[]{:T(512)S(6)} parameter(7) + %select_n.2213.clone.1 = pred[128,3,128,256]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.905), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.1273 = f32[3,128,128,256]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.2055.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_6.1273), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.1887 = f32[]{:T(128)} parameter(5) + %div.2638.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_5.1887), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2637.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.2055.clone.1, %div.2638.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2212.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2213.clone.1, %bitcast.2055.clone.1, %div.2637.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.6006.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4441.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.6006.clone.1), dimensions={}, metadata={op_name="broadcast.340"} + %mul.4780.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2212.clone.1, %broadcast.4441.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.673 = f32[128,3,128,256]{3,2,1,0:T(8,128)} parameter(8) + %constant.6010.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4781.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.6010.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4779.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.673, %mul.4781.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3656.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4780.clone.1, %mul.4779.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4166 = f32[]{:T(128)S(6)} parameter(2) + %div.2634.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4166), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.390.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2212.clone.1, %select_n.2212.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.6009.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4778.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.6009.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4776.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.390.clone.1, %mul.4778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.2169 = f32[128,3,128,256]{3,2,1,0:T(8,128)} parameter(4) + %constant.6008.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4777.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.6008.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4775.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2169, %mul.4777.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3655.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4776.clone.1, %mul.4775.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5190 = f32[]{:T(128)S(6)} parameter(1) + %div.2633.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5190), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2632.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3655.clone.1, %div.2633.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.162.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} sqrt(%div.2632.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.6007.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3654.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.6007.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3653.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.162.clone.1, %add.3654.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1421.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%div.2634.clone.1, %add.3653.clone.1), metadata={op_name="multiply.286"} + %div.2631.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3656.clone.1, %multiply.1421.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4773.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4434, %broadcast.4441.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3652.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} add(%div.2631.clone.1, %mul.4773.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4772.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.4774.clone.1, %add.3652.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3651.clone.1 = f32[128,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4434, %mul.4772.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.572 = f32[128,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3651.clone.1, %add.3651.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6207 = f32[]{:T(128)} constant(0) + %reduce.719 = f32[]{:T(128)} reduce(%square.572, %constant.6207), dimensions={0,1,2,3}, to_apply=%region_208.234, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.720.clone.1 = f32[]{:T(128)} reduce(%integer_pow.390.clone.1, %constant.6207), dimensions={0,1,2,3}, to_apply=%region_175.201, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.640 = (f32[]{:T(128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.719, %add.3651.clone.1, %add.3655.clone.1, %add.3656.clone.1, %reduce.720.clone.1) +} + +%region_162.188 (reduce_sum.412: f32[], reduce_sum.413: f32[]) -> f32[] { + %reduce_sum.412 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.413 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.414 = f32[]{:T(128)} add(%reduce_sum.412, %reduce_sum.413), 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.579 (param_0.4449: bf16[384,128,192]) -> f32[] { + %param_0.4449 = bf16[384,128,192]{1,0,2:T(8,128)(2,1)} parameter(0) + %broadcast_in_dim.1323 = f32[384,128,192]{1,0,2:T(8,128)} convert(%param_0.4449), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.1217 = f32[384,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1323), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.575 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.1217, %bitcast.1217), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6222 = f32[]{:T(128)} constant(0) + ROOT %reduce.721 = f32[]{:T(128)} reduce(%square.575, %constant.6222), dimensions={0,1,2,3}, to_apply=%region_162.188, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_228.254 (reduce_sum.672: f32[], reduce_sum.673: f32[]) -> f32[] { + %reduce_sum.672 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.673 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.1000 = f32[]{:T(128)} add(%reduce_sum.672, %reduce_sum.673), 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":[]}} +} + +%region_194.220 (reduce_sum.547: f32[], reduce_sum.548: f32[]) -> f32[] { + %reduce_sum.547 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.548 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.762 = f32[]{:T(128)} add(%reduce_sum.547, %reduce_sum.548), 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.584 (param_0.4417: f32[], param_1.5173: f32[], param_2.4149: f32[], param_3.2805: f32[384,1,128,192], param_4.2152: f32[384,1,128,192], param_5.1870: f32[], param_6.1256: bf16[384,128,192], param_7.888: pred[], param_8.657: f32[384,1,128,192]) -> (f32[], f32[384,1,128,192], f32[384,1,128,192], f32[384,1,128,192], f32[]) { + %param_3.2805 = f32[384,1,128,192]{2,3,1,0:T(8,128)} parameter(3) + %copy.1517.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} copy(%param_3.2805), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'moe_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\'].value"} + %param_2.4149 = f32[]{:T(128)S(6)} parameter(2) + %mul.4632.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%param_2.4149), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.888 = pred[]{:T(512)S(6)} parameter(7) + %select_n.2157.clone.1 = pred[384,1,128,192]{2,0,3,1:T(8,128)(4,1)} broadcast(%param_7.888), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.1256 = bf16[384,128,192]{1,0,2:T(8,128)(2,1)S(1)} parameter(6) + %broadcast_in_dim.1499.clone.1 = f32[384,128,192]{1,0,2:T(8,128)} convert(%param_6.1256), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/broadcast_in_dim" stack_frame_id=0} + %bitcast.2021.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1499.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.1870 = f32[]{:T(128)} parameter(5) + %div.2518.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%param_5.1870), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2517.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} divide(%bitcast.2021.clone.1, %div.2518.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2156.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} select(%select_n.2157.clone.1, %bitcast.2021.clone.1, %div.2517.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.5905.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4389.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%constant.5905.clone.1), dimensions={}, metadata={op_name="broadcast.2537"} + %mul.4638.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%select_n.2156.clone.1, %broadcast.4389.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.657 = f32[384,1,128,192]{2,3,1,0:T(8,128)} parameter(8) + %copy.1519.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} copy(%param_8.657), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'optimizer\'][\'opt_state\'][0][\'mu\'][\'decoder\'][\'moe_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\'].value"} + %constant.5909.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4639.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%constant.5909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4637.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%copy.1519.clone.1, %mul.4639.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3566.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} add(%mul.4638.clone.1, %mul.4637.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5173 = f32[]{:T(128)S(6)} parameter(1) + %div.2514.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%param_1.5173), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.374.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%select_n.2156.clone.1, %select_n.2156.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.5908.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4636.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%constant.5908.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4634.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%integer_pow.374.clone.1, %mul.4636.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.2152 = f32[384,1,128,192]{2,3,1,0:T(8,128)} parameter(4) + %copy.1518.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} copy(%param_4.2152), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'optimizer\'][\'opt_state\'][0][\'nu\'][\'decoder\'][\'moe_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\'].value"} + %constant.5907.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4635.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%constant.5907.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4633.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%copy.1518.clone.1, %mul.4635.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3565.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} add(%mul.4634.clone.1, %mul.4633.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_0.4417 = f32[]{:T(128)S(6)} parameter(0) + %div.2513.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%param_0.4417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2512.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} divide(%add.3565.clone.1, %div.2513.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.146.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} sqrt(%div.2512.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.5906.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3564.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} broadcast(%constant.5906.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3563.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} add(%sqrt.146.clone.1, %add.3564.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1405.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%div.2514.clone.1, %add.3563.clone.1), metadata={op_name="multiply.306"} + %div.2511.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} divide(%add.3566.clone.1, %multiply.1405.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4631.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%copy.1517.clone.1, %broadcast.4389.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3562.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} add(%div.2511.clone.1, %mul.4631.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4630.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%mul.4632.clone.1, %add.3562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3561.clone.1 = f32[384,1,128,192]{2,0,3,1:T(8,128)} add(%copy.1517.clone.1, %mul.4630.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.576 = f32[384,1,128,192]{2,0,3,1:T(8,128)} multiply(%add.3561.clone.1, %add.3561.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6190 = f32[]{:T(128)} constant(0) + %reduce.722 = f32[]{:T(128)} reduce(%square.576, %constant.6190), dimensions={0,1,2,3}, to_apply=%region_228.254, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.723.clone.1 = f32[]{:T(128)} reduce(%integer_pow.374.clone.1, %constant.6190), dimensions={0,1,2,3}, to_apply=%region_194.220, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.654 = (f32[]{:T(128)}, f32[384,1,128,192]{2,0,3,1:T(8,128)}, f32[384,1,128,192]{2,0,3,1:T(8,128)}, f32[384,1,128,192]{2,0,3,1:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.722, %add.3561.clone.1, %add.3565.clone.1, %add.3566.clone.1, %reduce.723.clone.1) +} + +%region_136.162 (reduce_sum.198: f32[], reduce_sum.204: f32[]) -> f32[] { + %reduce_sum.198 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.204 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.205 = f32[]{:T(128)} add(%reduce_sum.198, %reduce_sum.204), 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.597 (param_0.4470: f32[3,18432,128]) -> f32[] { + %param_0.4470 = f32[3,18432,128]{2,1,0:T(8,128)} parameter(0) + %bitcast.1236 = f32[18432,3,128]{2,0,1:T(8,128)} bitcast(%param_0.4470), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.579 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%bitcast.1236, %bitcast.1236), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6243 = f32[]{:T(128)} constant(0) + ROOT %reduce.724 = f32[]{:T(128)} reduce(%square.579, %constant.6243), dimensions={0,1,2}, to_apply=%region_136.162, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_135.161 (reduce_sum.192: f32[], reduce_sum.193: f32[]) -> f32[] { + %reduce_sum.192 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.193 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + 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.599 (param_0.4471: f32[3,128,18432]) -> f32[] { + %param_0.4471 = f32[3,128,18432]{2,1,0:T(8,128)S(1)} parameter(0) + %bitcast.1240 = f32[128,3,18432]{2,0,1:T(8,128)} bitcast(%param_0.4471), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.582 = f32[128,3,18432]{2,0,1:T(8,128)} multiply(%bitcast.1240, %bitcast.1240), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6244 = f32[]{:T(128)} constant(0) + ROOT %reduce.725 = f32[]{:T(128)} reduce(%square.582, %constant.6244), dimensions={0,1,2}, to_apply=%region_135.161, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_134.160 (reduce_sum.173: f32[], reduce_sum.181: f32[]) -> f32[] { + %reduce_sum.173 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.181 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.188 = f32[]{:T(128)} add(%reduce_sum.173, %reduce_sum.181), 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.601 (param_0.4472: f32[3,128,18432]) -> f32[] { + %param_0.4472 = f32[3,128,18432]{2,1,0:T(8,128)} parameter(0) + %bitcast.1244 = f32[128,3,18432]{2,0,1:T(8,128)} bitcast(%param_0.4472), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.585 = f32[128,3,18432]{2,0,1:T(8,128)} multiply(%bitcast.1244, %bitcast.1244), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6245 = f32[]{:T(128)} constant(0) + ROOT %reduce.726 = f32[]{:T(128)} reduce(%square.585, %constant.6245), dimensions={0,1,2}, to_apply=%region_134.160, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_201.227 (reduce_sum.573: f32[], reduce_sum.574: f32[]) -> f32[] { + %reduce_sum.573 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.574 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.811 = f32[]{:T(128)} add(%reduce_sum.573, %reduce_sum.574), 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":[]}} +} + +%region_168.194 (reduce_sum.436: f32[], reduce_sum.440: f32[]) -> f32[] { + %reduce_sum.436 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.440 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.441 = f32[]{:T(128)} add(%reduce_sum.436, %reduce_sum.440), 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.602 (param_0.4440: f32[18432,3,128], param_1.5196: f32[], param_2.4172: f32[], param_3.2828: f32[], param_4.2175: f32[18432,3,128], param_5.1893: f32[], param_6.1279: f32[3,18432,128], param_7.911: pred[], param_8.679: f32[18432,3,128]) -> (f32[], f32[18432,3,128], f32[18432,3,128], f32[18432,3,128], f32[]) { + %param_0.4440 = f32[18432,3,128]{2,0,1:T(8,128)} parameter(0) + %param_3.2828 = f32[]{:T(128)S(6)} parameter(3) + %mul.4821.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%param_3.2828), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.911 = pred[]{:T(512)S(6)} parameter(7) + %select_n.2233.clone.1 = pred[18432,3,128]{2,0,1:T(8,128)(4,1)} broadcast(%param_7.911), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.1279 = f32[3,18432,128]{2,1,0:T(8,128)S(1)} parameter(6) + %bitcast.2065.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} bitcast(%param_6.1279), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.1893 = f32[]{:T(128)} parameter(5) + %div.2678.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%param_5.1893), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2677.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} divide(%bitcast.2065.clone.1, %div.2678.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.2232.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} select(%select_n.2233.clone.1, %bitcast.2065.clone.1, %div.2677.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.6042.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.4455.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%constant.6042.clone.1), dimensions={}, metadata={op_name="broadcast.346"} + %mul.4827.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%select_n.2232.clone.1, %broadcast.4455.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.679 = f32[18432,3,128]{2,0,1:T(8,128)} parameter(8) + %constant.6046.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.4828.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%constant.6046.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4826.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%param_8.679, %mul.4828.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3685.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} add(%mul.4827.clone.1, %mul.4826.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.4172 = f32[]{:T(128)S(6)} parameter(2) + %div.2674.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%param_2.4172), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.395.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%select_n.2232.clone.1, %select_n.2232.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.6045.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.4825.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%constant.6045.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4823.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %mul.4825.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.2175 = f32[18432,3,128]{2,0,1:T(8,128)} parameter(4) + %constant.6044.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.4824.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%constant.6044.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4822.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%param_4.2175, %mul.4824.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3684.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} add(%mul.4823.clone.1, %mul.4822.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.5196 = f32[]{:T(128)S(6)} parameter(1) + %div.2673.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%param_1.5196), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2672.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} divide(%add.3684.clone.1, %div.2673.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.167.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} sqrt(%div.2672.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.6043.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.3683.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} broadcast(%constant.6043.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.3682.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} add(%sqrt.167.clone.1, %add.3683.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.1426.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%div.2674.clone.1, %add.3682.clone.1), metadata={op_name="multiply.279"} + %div.2671.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} divide(%add.3685.clone.1, %multiply.1426.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.4820.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%param_0.4440, %broadcast.4455.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3681.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} add(%div.2671.clone.1, %mul.4820.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.4819.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%mul.4821.clone.1, %add.3681.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.3680.clone.1 = f32[18432,3,128]{2,0,1:T(8,128)} add(%param_0.4440, %mul.4819.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.586 = f32[18432,3,128]{2,0,1:T(8,128)} multiply(%add.3680.clone.1, %add.3680.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.6213 = f32[]{:T(128)} constant(0) + %reduce.727 = f32[]{:T(128)} reduce(%square.586, %constant.6213), dimensions={0,1,2}, to_apply=%region_201.227, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.730.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.6213), dimensions={0,1,2}, to_apply=%region_168.194, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.645 = (f32[]{:T(128)}, f32[18432,3,128]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.727, %add.3680.clone.1, %add.3684.clone.1, %add.3685.clone.1, %reduce.730.clone.1) +} + +%region_200.226 (reduce_sum.568: f32[], reduce_sum.569: f32[]) -> f32[] { + %reduce_sum.568 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.569 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.804 = f32[]{:T(128)} add(%reduce_sum.568, %reduce_sum.569), 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":[]}} +} + +%region_167.193 (reduce_sum.433: f32[], reduce_sum.434: f32[]) -> f32[] { + %reduce_sum.433 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.434 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.435 = f32[]{:T(128)} add(%reduce_sum.433, %reduce_sum.434), 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":[]}} +} + +%region_166.192 (reduce_sum.427: f32[], reduce_sum.428: f32[]) -> f32[] { diff --git a/tests/utils/reference_hlo_llama3_8b.txt b/tests/utils/reference_hlo_llama3_8b.txt index 27c6529df2..3ee582b04f 100644 --- a/tests/utils/reference_hlo_llama3_8b.txt +++ b/tests/utils/reference_hlo_llama3_8b.txt @@ -1,4 +1,4 @@ -HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias) }, entry_computation_layout={(s32[]{:T(128)}, f32[4096]{0:T(1024)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, /*index=5*/f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, /*index=10*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, s32[]{:T(128)}, f32[4096]{0:T(1024)}, /*index=15*/f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, /*index=20*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, /*index=25*/f32[128256,4096]{1,0:T(8,128)}, f32[4096]{0:T(1024)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, /*index=30*/f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, /*index=35*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, s32[]{:T(128)}, s32[4,128]{1,0:T(4,128)}, /*index=40*/s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)})->(s32[]{:T(128)}, f32[4096]{0:T(1024)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, /*index=5*/f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, /*index=10*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, s32[]{:T(128)}, f32[4096]{0:T(1024)}, /*index=15*/f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, /*index=20*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, /*index=25*/f32[128256,4096]{1,0:T(8,128)}, f32[4096]{0:T(1024)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, /*index=30*/f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, /*index=35*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, s32[]{:T(128)}, f32[]{:T(128)}, /*index=40*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=45*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 +HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias), {42}: (42, {}, may-alias), {43}: (43, {}, may-alias), {44}: (44, {}, may-alias), {45}: (45, {}, may-alias), {46}: (46, {}, may-alias), {47}: (47, {}, may-alias), {48}: (48, {}, may-alias), {49}: (49, {}, may-alias), {50}: (50, {}, may-alias), {51}: (51, {}, may-alias), {52}: (52, {}, may-alias), {53}: (53, {}, may-alias), {54}: (54, {}, may-alias), {55}: (55, {}, may-alias), {56}: (56, {}, may-alias), {57}: (57, {}, may-alias), {58}: (58, {}, may-alias), {59}: (59, {}, may-alias), {60}: (60, {}, may-alias), {61}: (61, {}, may-alias), {62}: (62, {}, may-alias), {63}: (63, {}, may-alias), {64}: (64, {}, may-alias), {65}: (65, {}, may-alias), {66}: (66, {}, may-alias), {67}: (67, {}, may-alias), {68}: (68, {}, may-alias) }, entry_computation_layout={(f32[4096]{0:T(1024)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=5*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=10*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=15*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, /*index=20*/f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, u32[4]{0:T(128)}, /*index=25*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=30*/f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, /*index=35*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, /*index=40*/u32[4]{0:T(128)}, f32[128256,1024]{1,0:T(8,128)}, s32[]{:T(128)}, f32[4096]{0:T(1024)}, f32[1024,4,14336]{2,1,0:T(4,128)}, /*index=45*/f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, /*index=50*/f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, /*index=55*/f32[4096]{0:T(1024)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, /*index=60*/f32[4096,4]{0,1:T(4,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, /*index=65*/f32[1024,128256]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, s32[]{:T(128)}, u32[]{:T(128)}, s32[1,128]{1,0:T(1,128)}, /*index=70*/s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)})->(f32[4096]{0:T(1024)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=5*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=10*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=15*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, /*index=20*/f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, u32[4]{0:T(128)}, /*index=25*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=30*/f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, /*index=35*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, /*index=40*/u32[4]{0:T(128)}, f32[128256,1024]{1,0:T(8,128)}, s32[]{:T(128)}, f32[4096]{0:T(1024)}, f32[1024,4,14336]{2,1,0:T(4,128)}, /*index=45*/f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, /*index=50*/f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, /*index=55*/f32[4096]{0:T(1024)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, /*index=60*/f32[4096,4]{0,1:T(4,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, /*index=65*/f32[1024,128256]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, s32[]{:T(128)}, u32[]{:T(128)}, f32[]{:T(128)}, /*index=70*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=75*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 FileNames @@ -9,669 +9,621 @@ FileLocations StackFrames -%fused_computation (param_0.2: bf16[128256,4096], param_1.7: s32[1024]) -> bf16[512,4096] { - %param_0.2 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) +%region_0.1.clone (reduce_sum.416: s32[], reduce_sum.420: s32[]) -> s32[] { + %reduce_sum.416 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.420 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.421 = s32[]{:T(128)} add(%reduce_sum.416, %reduce_sum.420), 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"]}]}} +} + +%region_8.11.clone (dot_general.243: bf16[], dot_general.244: bf16[]) -> bf16[] { + %dot_general.243 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + %dot_general.244 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + ROOT %add.500 = bf16[]{:T(256)} add(%dot_general.243, %dot_general.244), metadata={op_name="add.51"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation (param_0.2: bf16[128256,1024], param_1.7: s32[1024]) -> bf16[512,1024] { + %param_0.2 = bf16[128256,1024]{1,0:T(8,128)(2,1)} parameter(0) %param_1.7 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.1 = s32[1024]{0:T(1024)} custom-call(%param_1.7), 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.6 = s32[512]{0:T(512)} slice(%custom-call.1), 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.342 = s32[4,128]{1,0:T(4,128)} reshape(%slice.6), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.326 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.342), 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.4 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} gather(%param_0.2, %transpose.326), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,4096}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %transpose.325 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} transpose(%gather.4), 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.341 = bf16[512,4096]{1,0:T(8,128)(2,1)} reshape(%transpose.325), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %custom-call.1 = s32[1024]{0:T(1024)} custom-call(%param_1.7), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %slice.34 = s32[512]{0:T(512)} slice(%custom-call.1), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %reshape.661 = s32[4,128]{1,0:T(4,128)} reshape(%slice.34), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.326 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.661), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %gather.4 = bf16[4,128,1024]{2,1,0:T(8,128)(2,1)} gather(%param_0.2, %transpose.326), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,1024}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.325 = bf16[4,128,1024]{2,1,0:T(8,128)(2,1)} transpose(%gather.4), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %reshape.660 = bf16[512,1024]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.325), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} } -%region_33.38.clone (scatter-add.6: bf16[], scatter-add.7: bf16[]) -> bf16[] { - %scatter-add.6 = 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.7 = 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.476 = bf16[]{:T(256)} add(%scatter-add.6, %scatter-add.7), 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":[]}} +%region_30.35.clone (scatter-add.6: bf16[], scatter-add.7: bf16[]) -> bf16[] { + %scatter-add.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add"} + %scatter-add.7 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add"} + ROOT %add.505 = bf16[]{:T(256)} add(%scatter-add.6, %scatter-add.7), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.1 (param_0.3: bf16[128256,4096], param_1.5: s32[512], param_2.4: bf16[512,4096]) -> bf16[128256,4096] { - %param_0.3 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) +%fused_computation.1 (param_0.3: bf16[128256,1024], param_1.5: s32[512], param_2.4: bf16[512,1024]) -> bf16[128256,1024] { + %param_0.3 = bf16[128256,1024]{1,0:T(8,128)(2,1)} parameter(0) %param_1.5 = s32[512]{0:T(512)S(1)} parameter(1) - %reshape.349 = 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.331 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.349), 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,4096]{1,0:T(8,128)(2,1)S(1)} parameter(2) - %reshape.350 = bf16[4,128,4096]{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.332 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} transpose(%reshape.350), 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[128256,4096]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.331, %transpose.332), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_33.38.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.668 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.5), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %transpose.331 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.668), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %param_2.4 = bf16[512,1024]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %reshape.669 = bf16[4,128,1024]{2,1,0:T(8,128)(2,1)} reshape(%param_2.4), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %transpose.332 = bf16[4,128,1024]{2,1,0:T(8,128)(2,1)} transpose(%reshape.669), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + ROOT %scatter.2 = bf16[128256,1024]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.331, %transpose.332), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_30.35.clone, metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} +} + +%add.15.clone (x.31: bf16[], y.31: bf16[]) -> bf16[] { + %x.31 = bf16[]{:T(256)} parameter(0) + %y.31 = bf16[]{:T(256)} parameter(1) + ROOT %add.504 = bf16[]{:T(256)} add(%x.31, %y.31), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_32.37 (reduce_sum.190: f32[], reduce_sum.191: f32[]) -> f32[] { +%all-reduce-scatter.7 (input.7: bf16[4096,128256]) -> bf16[1024,128256] { + %input.7 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(0) + %all-reduce.56 = bf16[4096,128256]{1,0:T(8,128)(2,1)} all-reduce(%input.7), channel_id=62, replica_groups={{0,1,2,3}}, use_global_device_ids=true, to_apply=%add.15.clone, frontend_attributes={from-cross-replica-sharding="true"}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"2"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} + %partition-id.24 = u32[] partition-id() + %constant.755 = u32[] constant(1024) + %multiply.258 = u32[]{:T(128)} multiply(%partition-id.24, %constant.755) + %constant.756 = u32[] constant(0) + ROOT %dynamic-slice.95 = bf16[1024,128256]{1,0:T(8,128)(2,1)} dynamic-slice(%all-reduce.56, %multiply.258, %constant.756), dynamic_slice_sizes={1024,128256}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294964223","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%region_29.34 (reduce_sum.190: f32[], reduce_sum.191: f32[]) -> f32[] { %reduce_sum.190 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.191 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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} - %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} - %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) -} - -%region_34.39 (reduce_sum.196: f32[], reduce_sum.197: f32[]) -> f32[] { +%fused_computation.258 (param_0.1282: bf16[1024,128256]) -> f32[] { + %param_0.1282 = bf16[1024,128256]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.845 = f32[1024,128256]{1,0:T(8,128)} convert(%param_0.1282), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %square.158 = f32[1024,128256]{1,0:T(8,128)} multiply(%convert_element_type.845, %convert_element_type.845), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1270 = f32[]{:T(128)} constant(0) + ROOT %reduce.131 = f32[]{:T(128)} reduce(%square.158, %constant.1270), dimensions={0,1}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_31.36 (reduce_sum.196: f32[], reduce_sum.197: f32[]) -> f32[] { %reduce_sum.196 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.197 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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} +%fused_computation.259 (param_0.1281: bf16[128256,1024]) -> f32[] { + %param_0.1281 = bf16[128256,1024]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.847 = f32[128256,1024]{1,0:T(8,128)} convert(%param_0.1281), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %square.160 = f32[128256,1024]{1,0:T(8,128)} multiply(%convert_element_type.847, %convert_element_type.847), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1269 = f32[]{:T(128)} constant(0) + ROOT %reduce.132 = f32[]{:T(128)} reduce(%square.160, %constant.1269), dimensions={0,1}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%region_60.65 (reduce_sum.338: f32[], reduce_sum.339: f32[]) -> f32[] { +%region_57.62 (reduce_sum.338: f32[], reduce_sum.339: f32[]) -> f32[] { %reduce_sum.338 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.339 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.329 = 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":[]}} } -%region_46.51 (reduce_sum.259: f32[], reduce_sum.260: f32[]) -> f32[] { +%region_43.48 (reduce_sum.259: f32[], reduce_sum.260: f32[]) -> f32[] { %reduce_sum.259 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.260 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %add.774.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.908.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - %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.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) -} - -%region_59.64 (reduce_sum.331: f32[], reduce_sum.332: f32[]) -> f32[] { +%fused_computation.260 (param_0.1267: f32[128256,1024], param_1.1409: f32[], param_2.1093: f32[], param_3.735: f32[], param_4.452: f32[128256,1024], param_5.398: f32[], param_6.253: bf16[128256,1024], param_7.169: pred[], param_8.98: f32[128256,1024]) -> (f32[], f32[128256,1024], f32[128256,1024], f32[128256,1024], f32[]) { + %param_0.1267 = f32[128256,1024]{1,0:T(8,128)} parameter(0) + %param_3.735 = f32[]{:T(128)S(6)} parameter(3) + %mul.1570.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%param_3.735), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.169 = pred[]{:T(512)S(6)} parameter(7) + %select_n.254.clone.1 = pred[128256,1024]{1,0:T(8,128)(4,1)} broadcast(%param_7.169), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.253 = bf16[128256,1024]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.954.clone.1 = f32[128256,1024]{1,0:T(8,128)} convert(%param_6.253), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_5.398 = f32[]{:T(128)} parameter(5) + %div.763.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%param_5.398), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.762.clone.1 = f32[128256,1024]{1,0:T(8,128)} divide(%convert_element_type.954.clone.1, %div.763.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.253.clone.1 = f32[128256,1024]{1,0:T(8,128)} select(%select_n.254.clone.1, %convert_element_type.954.clone.1, %div.762.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1198.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.659.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%constant.1198.clone.1), dimensions={}, metadata={op_name="broadcast.63"} + %mul.1576.clone.1 = f32[128256,1024]{1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.659.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.98 = f32[128256,1024]{1,0:T(8,128)} parameter(8) + %constant.1202.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.1577.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%constant.1202.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1575.clone.1 = f32[128256,1024]{1,0:T(8,128)} multiply(%param_8.98, %mul.1577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.820.clone.1 = f32[128256,1024]{1,0:T(8,128)} add(%mul.1576.clone.1, %mul.1575.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1093 = f32[]{:T(128)S(6)} parameter(2) + %div.759.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%param_2.1093), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.63.clone.1 = f32[128256,1024]{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.1201.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.1574.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%constant.1201.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1572.clone.1 = f32[128256,1024]{1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1574.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.452 = f32[128256,1024]{1,0:T(8,128)} parameter(4) + %constant.1200.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.1573.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%constant.1200.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1571.clone.1 = f32[128256,1024]{1,0:T(8,128)} multiply(%param_4.452, %mul.1573.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.819.clone.1 = f32[128256,1024]{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_1.1409 = f32[]{:T(128)S(6)} parameter(1) + %div.758.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%param_1.1409), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.757.clone.1 = f32[128256,1024]{1,0:T(8,128)} divide(%add.819.clone.1, %div.758.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.61.clone.1 = f32[128256,1024]{1,0:T(8,128)} sqrt(%div.757.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1199.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.818.clone.1 = f32[128256,1024]{1,0:T(8,128)} broadcast(%constant.1199.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.817.clone.1 = f32[128256,1024]{1,0:T(8,128)} add(%sqrt.61.clone.1, %add.818.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.280.clone.1 = f32[128256,1024]{1,0:T(8,128)} multiply(%div.759.clone.1, %add.817.clone.1), metadata={op_name="multiply.41"} + %div.756.clone.1 = f32[128256,1024]{1,0:T(8,128)} divide(%add.820.clone.1, %multiply.280.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1569.clone.1 = f32[128256,1024]{1,0:T(8,128)} multiply(%param_0.1267, %broadcast.659.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.816.clone.1 = f32[128256,1024]{1,0:T(8,128)} add(%div.756.clone.1, %mul.1569.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1568.clone.1 = f32[128256,1024]{1,0:T(8,128)} multiply(%mul.1570.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.815.clone.1 = f32[128256,1024]{1,0:T(8,128)} add(%param_0.1267, %mul.1568.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.161 = f32[128256,1024]{1,0:T(8,128)} multiply(%add.815.clone.1, %add.815.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1255 = f32[]{:T(128)} constant(0) + %reduce.133 = f32[]{:T(128)} reduce(%square.161, %constant.1255), dimensions={0,1}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.135.clone.1 = f32[]{:T(128)} reduce(%integer_pow.63.clone.1, %constant.1255), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.134 = (f32[]{:T(128)}, f32[128256,1024]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.133, %add.815.clone.1, %add.819.clone.1, %add.820.clone.1, %reduce.135.clone.1) +} + +%region_56.61 (reduce_sum.331: f32[], reduce_sum.332: f32[]) -> f32[] { %reduce_sum.331 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.332 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.323 = f32[]{:T(128)} add(%reduce_sum.331, %reduce_sum.332), 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":[]}} } -%region_45.50 (reduce_sum.253: f32[], reduce_sum.254: f32[]) -> f32[] { +%region_42.47 (reduce_sum.253: f32[], reduce_sum.254: f32[]) -> f32[] { %reduce_sum.253 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.254 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %add.780.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.914.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - %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.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) -} - -%region_25.30 (reduce_sum.154: f32[], reduce_sum.155: f32[]) -> f32[] { +%fused_computation.261 (param_0.1268: f32[1024,128256], param_1.1410: f32[], param_2.1094: f32[], param_3.736: f32[], param_4.453: f32[1024,128256], param_5.399: f32[], param_6.254: bf16[1024,128256], param_7.170: pred[], param_8.99: f32[1024,128256]) -> (f32[], f32[1024,128256], f32[1024,128256], f32[1024,128256], f32[]) { + %param_0.1268 = f32[1024,128256]{1,0:T(8,128)} parameter(0) + %param_3.736 = f32[]{:T(128)S(6)} parameter(3) + %mul.1580.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%param_3.736), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.170 = pred[]{:T(512)S(6)} parameter(7) + %select_n.258.clone.1 = pred[1024,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.170), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.254 = bf16[1024,128256]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.956.clone.1 = f32[1024,128256]{1,0:T(8,128)} convert(%param_6.254), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_5.399 = f32[]{:T(128)} parameter(5) + %div.771.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%param_5.399), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.770.clone.1 = f32[1024,128256]{1,0:T(8,128)} divide(%convert_element_type.956.clone.1, %div.771.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.257.clone.1 = f32[1024,128256]{1,0:T(8,128)} select(%select_n.258.clone.1, %convert_element_type.956.clone.1, %div.770.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1204.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.661.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%constant.1204.clone.1), dimensions={}, metadata={op_name="broadcast.64"} + %mul.1586.clone.1 = f32[1024,128256]{1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.661.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.99 = f32[1024,128256]{1,0:T(8,128)} parameter(8) + %constant.1208.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.1587.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%constant.1208.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1585.clone.1 = f32[1024,128256]{1,0:T(8,128)} multiply(%param_8.99, %mul.1587.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.826.clone.1 = f32[1024,128256]{1,0:T(8,128)} add(%mul.1586.clone.1, %mul.1585.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1094 = f32[]{:T(128)S(6)} parameter(2) + %div.767.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%param_2.1094), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.64.clone.1 = f32[1024,128256]{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.1207.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.1584.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%constant.1207.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1582.clone.1 = f32[1024,128256]{1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.453 = f32[1024,128256]{1,0:T(8,128)} parameter(4) + %constant.1206.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.1583.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%constant.1206.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1581.clone.1 = f32[1024,128256]{1,0:T(8,128)} multiply(%param_4.453, %mul.1583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.825.clone.1 = f32[1024,128256]{1,0:T(8,128)} add(%mul.1582.clone.1, %mul.1581.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1410 = f32[]{:T(128)S(6)} parameter(1) + %div.766.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%param_1.1410), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.765.clone.1 = f32[1024,128256]{1,0:T(8,128)} divide(%add.825.clone.1, %div.766.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.62.clone.1 = f32[1024,128256]{1,0:T(8,128)} sqrt(%div.765.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1205.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.824.clone.1 = f32[1024,128256]{1,0:T(8,128)} broadcast(%constant.1205.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.823.clone.1 = f32[1024,128256]{1,0:T(8,128)} add(%sqrt.62.clone.1, %add.824.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.281.clone.1 = f32[1024,128256]{1,0:T(8,128)} multiply(%div.767.clone.1, %add.823.clone.1), metadata={op_name="multiply.40"} + %div.764.clone.1 = f32[1024,128256]{1,0:T(8,128)} divide(%add.826.clone.1, %multiply.281.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1579.clone.1 = f32[1024,128256]{1,0:T(8,128)} multiply(%param_0.1268, %broadcast.661.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.822.clone.1 = f32[1024,128256]{1,0:T(8,128)} add(%div.764.clone.1, %mul.1579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1578.clone.1 = f32[1024,128256]{1,0:T(8,128)} multiply(%mul.1580.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.821.clone.1 = f32[1024,128256]{1,0:T(8,128)} add(%param_0.1268, %mul.1578.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.162 = f32[1024,128256]{1,0:T(8,128)} multiply(%add.821.clone.1, %add.821.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1256 = f32[]{:T(128)} constant(0) + %reduce.134 = f32[]{:T(128)} reduce(%square.162, %constant.1256), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.136.clone.1 = f32[]{:T(128)} reduce(%integer_pow.64.clone.1, %constant.1256), dimensions={0,1}, to_apply=%region_42.47, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.135 = (f32[]{:T(128)}, f32[1024,128256]{1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.134, %add.821.clone.1, %add.825.clone.1, %add.826.clone.1, %reduce.136.clone.1) +} + +%region_22.27 (reduce_sum.154: f32[], reduce_sum.155: f32[]) -> f32[] { %reduce_sum.154 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.155 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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} +%fused_computation.274 (param_0.1286: f32[4,14336,1024]) -> f32[] { + %param_0.1286 = f32[4,14336,1024]{2,0,1:T(4,128)} parameter(0) + %bitcast.467 = f32[14336,4,1024]{2,1,0:T(4,128)} bitcast(%param_0.1286), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.165 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%bitcast.467, %bitcast.467), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1274 = f32[]{:T(128)} constant(0) + ROOT %reduce.137 = f32[]{:T(128)} reduce(%square.165, %constant.1274), dimensions={0,1,2}, to_apply=%region_22.27, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%region_24.29 (reduce_sum.148: f32[], reduce_sum.149: f32[]) -> f32[] { +%region_21.26 (reduce_sum.148: f32[], reduce_sum.149: f32[]) -> f32[] { %reduce_sum.148 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.149 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.150 = f32[]{:T(128)} add(%reduce_sum.148, %reduce_sum.149), 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":[]}} } -%region_23.28 (reduce_sum.142: f32[], reduce_sum.143: f32[]) -> f32[] { +%region_20.25 (reduce_sum.142: f32[], reduce_sum.143: f32[]) -> f32[] { %reduce_sum.142 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.143 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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} - ROOT %tuple.155 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.125, %reduce.126.clone.1) +%fused_computation.276 (param_0.1287: f32[4,1024,14336], param_1.1422: f32[4,1024,14336]) -> (f32[], f32[]) { + %param_0.1287 = f32[4,1024,14336]{2,0,1:T(4,128)} parameter(0) + %bitcast.471 = f32[1024,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1287), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.168 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.471, %bitcast.471), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1275 = f32[]{:T(128)} constant(0) + %reduce.138 = f32[]{:T(128)} reduce(%square.168, %constant.1275), dimensions={0,1,2}, to_apply=%region_21.26, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1422 = f32[4,1024,14336]{2,0,1:T(4,128)} parameter(1) + %bitcast.475.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1422), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.171.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.475.clone.1, %bitcast.475.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.139.clone.1 = f32[]{:T(128)} reduce(%square.171.clone.1, %constant.1275), dimensions={0,1,2}, to_apply=%region_20.25, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.154 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.138, %reduce.139.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.279 (param_0.807: f32[14336,4,1024]) -> bf16[4,14336,1024] { + %param_0.807 = f32[14336,4,1024]{2,1,0:T(4,128)} parameter(0) + %copy.235 = bf16[14336,4,1024]{2,0,1:T(8,128)(2,1)} copy(%param_0.807), sharding={devices=[1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\'].value"} + ROOT %bitcast.476 = bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp()/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.280 (param_0.809: f32[1024,4,14336]) -> bf16[4,1024,14336] { + %param_0.809 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.236 = bf16[1024,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.809), sharding={devices=[4,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\'].value"} + ROOT %bitcast.477 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp()/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.281 (param_0.811: f32[1024,4,14336]) -> bf16[4,1024,14336] { + %param_0.811 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.237 = bf16[1024,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.811), sharding={devices=[4,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\'].value"} + ROOT %bitcast.478 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} } -%region_52.57 (reduce_sum.289: f32[], reduce_sum.290: f32[]) -> f32[] { +%region_49.54 (reduce_sum.289: f32[], reduce_sum.290: f32[]) -> f32[] { %reduce_sum.289 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.290 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.294 = f32[]{:T(128)} add(%reduce_sum.289, %reduce_sum.290), 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":[]}} } -%region_38.43 (reduce_sum.217: f32[], reduce_sum.218: f32[]) -> f32[] { +%region_35.40 (reduce_sum.217: f32[], reduce_sum.218: f32[]) -> f32[] { %reduce_sum.217 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.218 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %add.818.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.956.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - %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.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) -} - -%region_51.56 (reduce_sum.283: f32[], reduce_sum.287: f32[]) -> f32[] { +%fused_computation.282 (param_0.1275: f32[14336,4,1024], param_1.1417: f32[], param_2.1101: f32[], param_3.743: f32[], param_4.460: f32[14336,4,1024], param_5.406: f32[], param_6.261: f32[4,14336,1024], param_7.177: pred[], param_8.106: f32[14336,4,1024]) -> (f32[], f32[14336,4,1024], f32[14336,4,1024], f32[14336,4,1024], f32[]) { + %param_0.1275 = f32[14336,4,1024]{2,1,0:T(4,128)} parameter(0) + %param_3.743 = f32[]{:T(128)S(6)} parameter(3) + %mul.1624.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%param_3.743), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.177 = pred[]{:T(512)S(6)} parameter(7) + %select_n.278.clone.1 = pred[14336,4,1024]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.177), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.261 = f32[4,14336,1024]{2,0,1:T(4,128)} parameter(6) + %bitcast.623.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} bitcast(%param_6.261), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.406 = f32[]{:T(128)} parameter(5) + %div.811.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%param_5.406), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.810.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} divide(%bitcast.623.clone.1, %div.811.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.277.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.623.clone.1, %div.810.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1234.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.679.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%constant.1234.clone.1), dimensions={}, metadata={op_name="broadcast.71"} + %mul.1630.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.679.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.106 = f32[14336,4,1024]{2,1,0:T(4,128)} parameter(8) + %constant.1238.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.1631.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%constant.1238.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1629.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%param_8.106, %mul.1631.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.854.clone.1 = f32[14336,4,1024]{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.1101 = f32[]{:T(128)S(6)} parameter(2) + %div.807.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.69.clone.1 = f32[14336,4,1024]{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.1237.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.1628.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%constant.1237.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1626.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %mul.1628.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.460 = f32[14336,4,1024]{2,1,0:T(4,128)} parameter(4) + %constant.1236.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.1627.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%constant.1236.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1625.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%param_4.460, %mul.1627.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.853.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} add(%mul.1626.clone.1, %mul.1625.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1417 = f32[]{:T(128)S(6)} parameter(1) + %div.806.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%param_1.1417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.805.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} divide(%add.853.clone.1, %div.806.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.67.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} sqrt(%div.805.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1235.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.852.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} broadcast(%constant.1235.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.851.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} add(%sqrt.67.clone.1, %add.852.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.286.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%div.807.clone.1, %add.851.clone.1), metadata={op_name="multiply.33"} + %div.804.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} divide(%add.854.clone.1, %multiply.286.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1623.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%param_0.1275, %broadcast.679.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.850.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} add(%div.804.clone.1, %mul.1623.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1622.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%mul.1624.clone.1, %add.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.849.clone.1 = f32[14336,4,1024]{2,1,0:T(4,128)} add(%param_0.1275, %mul.1622.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.172 = f32[14336,4,1024]{2,1,0:T(4,128)} multiply(%add.849.clone.1, %add.849.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1263 = f32[]{:T(128)} constant(0) + %reduce.140 = f32[]{:T(128)} reduce(%square.172, %constant.1263), dimensions={0,1,2}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.143.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1263), dimensions={0,1,2}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.136 = (f32[]{:T(128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.140, %add.849.clone.1, %add.853.clone.1, %add.854.clone.1, %reduce.143.clone.1) +} + +%region_48.53 (reduce_sum.283: f32[], reduce_sum.287: f32[]) -> f32[] { %reduce_sum.283 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.287 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.288 = f32[]{:T(128)} add(%reduce_sum.283, %reduce_sum.287), 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":[]}} } -%region_37.42 (reduce_sum.211: f32[], reduce_sum.212: f32[]) -> f32[] { +%region_34.39 (reduce_sum.211: f32[], reduce_sum.212: f32[]) -> f32[] { %reduce_sum.211 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.212 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %broadcast.587.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.962.clone.1), dimensions={}, metadata={op_name="broadcast.54"} - %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} - %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.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) -} - -%region_50.55 (reduce_sum.280: f32[], reduce_sum.281: f32[]) -> f32[] { +%fused_computation.283 (param_0.1276: f32[1024,4,14336], param_1.1418: f32[], param_2.1102: f32[], param_3.744: f32[], param_4.461: f32[1024,4,14336], param_5.407: f32[], param_6.262: f32[4,1024,14336], param_7.178: pred[], param_8.107: f32[1024,4,14336]) -> (f32[], f32[1024,4,14336], f32[1024,4,14336], f32[1024,4,14336], f32[]) { + %param_0.1276 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.744 = f32[]{:T(128)S(6)} parameter(3) + %mul.1634.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.744), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.178 = pred[]{:T(512)S(6)} parameter(7) + %select_n.282.clone.1 = pred[1024,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.178), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.262 = f32[4,1024,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.625.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.262), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.407 = f32[]{:T(128)} parameter(5) + %div.819.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.407), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.818.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} divide(%bitcast.625.clone.1, %div.819.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.281.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.625.clone.1, %div.818.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1240.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.685.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1240.clone.1), dimensions={}, metadata={op_name="broadcast.73"} + %mul.1638.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.685.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.107 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(8) + %constant.1244.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.684.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1244.clone.1), dimensions={}, metadata={op_name="broadcast.72"} + %mul.1637.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%param_8.107, %broadcast.684.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.859.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%mul.1638.clone.1, %mul.1637.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) + %div.815.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.70.clone.1 = f32[1024,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.1243.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.683.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1243.clone.1), dimensions={}, metadata={op_name="broadcast.62"} + %mul.1636.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.683.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.461 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(4) + %constant.1242.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.682.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1242.clone.1), dimensions={}, metadata={op_name="broadcast.61"} + %mul.1635.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%param_4.461, %broadcast.682.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.858.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%mul.1636.clone.1, %mul.1635.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1418 = f32[]{:T(128)S(6)} parameter(1) + %div.814.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.813.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} divide(%add.858.clone.1, %div.814.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.68.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} sqrt(%div.813.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1241.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.680.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1241.clone.1), dimensions={}, metadata={op_name="broadcast.56"} + %add.857.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%sqrt.68.clone.1, %broadcast.680.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.287.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%div.815.clone.1, %add.857.clone.1), metadata={op_name="multiply.32"} + %div.812.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} divide(%add.859.clone.1, %multiply.287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1633.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1276, %broadcast.685.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.856.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%div.812.clone.1, %mul.1633.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1632.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%mul.1634.clone.1, %add.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.855.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%param_0.1276, %mul.1632.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.173 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%add.855.clone.1, %add.855.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1264 = f32[]{:T(128)} constant(0) + %reduce.141 = f32[]{:T(128)} reduce(%square.173, %constant.1264), dimensions={0,1,2}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.144.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1264), dimensions={0,1,2}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.137 = (f32[]{:T(128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.141, %add.855.clone.1, %add.858.clone.1, %add.859.clone.1, %reduce.144.clone.1) +} + +%region_47.52 (reduce_sum.280: f32[], reduce_sum.281: f32[]) -> f32[] { %reduce_sum.280 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.281 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.282 = f32[]{:T(128)} add(%reduce_sum.280, %reduce_sum.281), 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":[]}} } -%region_36.41 (reduce_sum.205: f32[], reduce_sum.206: f32[]) -> f32[] { +%region_33.38 (reduce_sum.205: f32[], reduce_sum.206: f32[]) -> f32[] { %reduce_sum.205 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.206 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %broadcast.593.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.968.clone.1), dimensions={}, metadata={op_name="broadcast.54"} - %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} - %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.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) -} - -%region_30.35 (reduce_sum.178: f32[], reduce_sum.182: f32[]) -> f32[] { +%fused_computation.284 (param_0.1277: f32[1024,4,14336], param_1.1419: f32[], param_2.1103: f32[], param_3.745: f32[], param_4.462: f32[1024,4,14336], param_5.408: f32[], param_6.263: f32[4,1024,14336], param_7.179: pred[], param_8.108: f32[1024,4,14336]) -> (f32[], f32[1024,4,14336], f32[1024,4,14336], f32[1024,4,14336], f32[]) { + %param_0.1277 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(0) + %param_3.745 = f32[]{:T(128)S(6)} parameter(3) + %mul.1641.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.745), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.179 = pred[]{:T(512)S(6)} parameter(7) + %select_n.286.clone.1 = pred[1024,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.179), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.263 = f32[4,1024,14336]{2,0,1:T(4,128)} parameter(6) + %bitcast.627.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.263), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.408 = f32[]{:T(128)} parameter(5) + %div.827.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.408), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.826.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} divide(%bitcast.627.clone.1, %div.827.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.285.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} select(%select_n.286.clone.1, %bitcast.627.clone.1, %div.826.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1246.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.691.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1246.clone.1), dimensions={}, metadata={op_name="broadcast.73"} + %mul.1645.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%select_n.285.clone.1, %broadcast.691.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.108 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(8) + %constant.1250.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.690.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1250.clone.1), dimensions={}, metadata={op_name="broadcast.72"} + %mul.1644.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%param_8.108, %broadcast.690.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.864.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%mul.1645.clone.1, %mul.1644.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1103 = f32[]{:T(128)S(6)} parameter(2) + %div.823.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1103), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.71.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%select_n.285.clone.1, %select_n.285.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.1249.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.689.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1249.clone.1), dimensions={}, metadata={op_name="broadcast.62"} + %mul.1643.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.71.clone.1, %broadcast.689.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.462 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(4) + %constant.1248.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.688.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1248.clone.1), dimensions={}, metadata={op_name="broadcast.61"} + %mul.1642.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%param_4.462, %broadcast.688.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.863.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%mul.1643.clone.1, %mul.1642.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1419 = f32[]{:T(128)S(6)} parameter(1) + %div.822.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1419), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.821.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} divide(%add.863.clone.1, %div.822.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.69.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} sqrt(%div.821.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1247.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.686.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} broadcast(%constant.1247.clone.1), dimensions={}, metadata={op_name="broadcast.56"} + %add.862.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%sqrt.69.clone.1, %broadcast.686.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.288.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%div.823.clone.1, %add.862.clone.1), metadata={op_name="multiply.31"} + %div.820.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} divide(%add.864.clone.1, %multiply.288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1640.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1277, %broadcast.691.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.861.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%div.820.clone.1, %mul.1640.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1639.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%mul.1641.clone.1, %add.861.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.860.clone.1 = f32[1024,4,14336]{2,1,0:T(4,128)} add(%param_0.1277, %mul.1639.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.174 = f32[1024,4,14336]{2,1,0:T(4,128)} multiply(%add.860.clone.1, %add.860.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1265 = f32[]{:T(128)} constant(0) + %reduce.142 = f32[]{:T(128)} reduce(%square.174, %constant.1265), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.145.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1265), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.138 = (f32[]{:T(128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.142, %add.860.clone.1, %add.863.clone.1, %add.864.clone.1, %reduce.145.clone.1) +} + +%region_27.32 (reduce_sum.178: f32[], reduce_sum.182: f32[]) -> f32[] { %reduce_sum.178 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.182 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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} +%fused_computation.300 (param_0.1283: f32[4,1024,32,128]) -> f32[] { + %param_0.1283 = f32[4,1024,32,128]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.482 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1283), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.177 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.482, %bitcast.482), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1271 = f32[]{:T(128)} constant(0) + ROOT %reduce.146 = f32[]{:T(128)} reduce(%square.177, %constant.1271), dimensions={0,1,2,3}, to_apply=%region_27.32, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%region_29.34 (reduce_sum.175: f32[], reduce_sum.176: f32[]) -> f32[] { +%region_26.31 (reduce_sum.175: f32[], reduce_sum.176: f32[]) -> f32[] { %reduce_sum.175 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.176 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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.302 (param_0.1284: f32[4,32,128,1024]) -> f32[] { + %param_0.1284 = f32[4,32,128,1024]{3,2,0,1:T(8,128)} parameter(0) + %bitcast.486 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} bitcast(%param_0.1284), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.180 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%bitcast.486, %bitcast.486), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1272 = f32[]{:T(128)} constant(0) + ROOT %reduce.147 = f32[]{:T(128)} reduce(%square.180, %constant.1272), dimensions={0,1,2,3}, to_apply=%region_26.31, 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.303 (param_0.855: f32[32,4,128,1024]) -> bf16[4,32,128,1024] { + %param_0.855 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} parameter(0) + %copy.238 = bf16[32,4,128,1024]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.855), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\'].value"} + ROOT %bitcast.487 = bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} } -%region_57.62 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { +%region_54.59 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { %reduce_sum.317 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.318 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.316 = 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":[]}} } -%region_43.48 (reduce_sum.241: f32[], reduce_sum.245: f32[]) -> f32[] { +%region_40.45 (reduce_sum.241: f32[], reduce_sum.245: f32[]) -> f32[] { %reduce_sum.241 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.245 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %add.791.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.926.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - %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.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) -} - -%region_56.61 (reduce_sum.310: f32[], reduce_sum.311: f32[]) -> f32[] { +%fused_computation.304 (param_0.1270: f32[1024,4,32,128], param_1.1412: f32[], param_2.1096: f32[], param_3.738: f32[], param_4.455: f32[1024,4,32,128], param_5.401: f32[], param_6.256: f32[4,1024,32,128], param_7.172: pred[], param_8.101: f32[1024,4,32,128]) -> (f32[], f32[1024,4,32,128], f32[1024,4,32,128], f32[1024,4,32,128], f32[]) { + %param_0.1270 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.738 = f32[]{:T(128)S(6)} parameter(3) + %mul.1597.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.738), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.172 = pred[]{:T(512)S(6)} parameter(7) + %select_n.266.clone.1 = pred[1024,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.172), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.256 = f32[4,1024,32,128]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.617.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.256), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.401 = f32[]{:T(128)} parameter(5) + %div.787.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.401), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.786.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.617.clone.1, %div.787.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.265.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.266.clone.1, %bitcast.617.clone.1, %div.786.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1216.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.669.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.1216.clone.1), dimensions={}, metadata={op_name="broadcast.65"} + %mul.1603.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.265.clone.1, %broadcast.669.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.101 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} parameter(8) + %constant.1220.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.1604.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.1220.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1602.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.101, %mul.1604.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.837.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1603.clone.1, %mul.1602.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1096 = f32[]{:T(128)S(6)} parameter(2) + %div.783.clone.1 = f32[1024,4,32,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.66.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,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.1219.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.1601.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.1219.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1599.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.66.clone.1, %mul.1601.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.455 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} parameter(4) + %constant.1218.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.1600.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.1218.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1598.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.455, %mul.1600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.836.clone.1 = f32[1024,4,32,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_1.1412 = f32[]{:T(128)S(6)} parameter(1) + %div.782.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1412), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.781.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} divide(%add.836.clone.1, %div.782.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.64.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} sqrt(%div.781.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1217.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.835.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.1217.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.834.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} add(%sqrt.64.clone.1, %add.835.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.283.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%div.783.clone.1, %add.834.clone.1), metadata={op_name="multiply.38"} + %div.780.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} divide(%add.837.clone.1, %multiply.283.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1596.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1270, %broadcast.669.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.833.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} add(%div.780.clone.1, %mul.1596.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1595.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1597.clone.1, %add.833.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.832.clone.1 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1270, %mul.1595.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.181 = f32[1024,4,32,128]{3,2,1,0:T(8,128)} multiply(%add.832.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1258 = f32[]{:T(128)} constant(0) + %reduce.148 = f32[]{:T(128)} reduce(%square.181, %constant.1258), dimensions={0,1,2,3}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.152.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1258), dimensions={0,1,2,3}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.139 = (f32[]{:T(128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.148, %add.832.clone.1, %add.836.clone.1, %add.837.clone.1, %reduce.152.clone.1) +} + +%region_53.58 (reduce_sum.310: f32[], reduce_sum.311: f32[]) -> f32[] { %reduce_sum.310 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.311 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.315 = f32[]{:T(128)} add(%reduce_sum.310, %reduce_sum.311), 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":[]}} } -%region_42.47 (reduce_sum.238: f32[], reduce_sum.239: f32[]) -> f32[] { +%region_39.44 (reduce_sum.238: f32[], reduce_sum.239: f32[]) -> f32[] { %reduce_sum.238 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.239 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %add.797.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.932.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - %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.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) -} - -%region_47.52 (reduce_sum.262: f32[], reduce_sum.266: f32[]) -> f32[] { - %reduce_sum.262 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.266 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - 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} - %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} - %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} +%fused_computation.305 (param_0.1271: f32[32,4,128,1024], param_1.1413: f32[], param_2.1097: f32[], param_3.739: f32[], param_4.456: f32[32,4,128,1024], param_5.402: f32[], param_6.257: f32[4,32,128,1024], param_7.173: pred[], param_8.102: f32[32,4,128,1024]) -> (f32[], f32[32,4,128,1024], f32[32,4,128,1024], f32[32,4,128,1024], f32[]) { + %param_0.1271 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} parameter(0) + %param_3.739 = f32[]{:T(128)S(6)} parameter(3) + %mul.1607.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%param_3.739), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.173 = pred[]{:T(512)S(6)} parameter(7) + %select_n.270.clone.1 = pred[32,4,128,1024]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.173), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.257 = f32[4,32,128,1024]{3,2,0,1:T(8,128)} parameter(6) + %bitcast.619.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} bitcast(%param_6.257), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.402 = f32[]{:T(128)} parameter(5) + %div.795.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%param_5.402), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.794.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} divide(%bitcast.619.clone.1, %div.795.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.269.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} select(%select_n.270.clone.1, %bitcast.619.clone.1, %div.794.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1222.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.671.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%constant.1222.clone.1), dimensions={}, metadata={op_name="broadcast.66"} + %mul.1613.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%select_n.269.clone.1, %broadcast.671.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.102 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} parameter(8) + %constant.1226.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.1614.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%constant.1226.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1612.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%param_8.102, %mul.1614.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.843.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} add(%mul.1613.clone.1, %mul.1612.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1097 = f32[]{:T(128)S(6)} parameter(2) + %div.791.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%param_2.1097), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.67.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,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.1225.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.1611.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%constant.1225.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1609.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%integer_pow.67.clone.1, %mul.1611.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.456 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} parameter(4) + %constant.1224.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.1610.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%constant.1224.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1608.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%param_4.456, %mul.1610.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.842.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} add(%mul.1609.clone.1, %mul.1608.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1413 = f32[]{:T(128)S(6)} parameter(1) + %div.790.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%param_1.1413), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.789.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} divide(%add.842.clone.1, %div.790.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.65.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} sqrt(%div.789.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1223.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.841.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} broadcast(%constant.1223.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.840.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} add(%sqrt.65.clone.1, %add.841.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.284.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%div.791.clone.1, %add.840.clone.1), metadata={op_name="multiply.37"} + %div.788.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} divide(%add.843.clone.1, %multiply.284.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1606.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%param_0.1271, %broadcast.671.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.839.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} add(%div.788.clone.1, %mul.1606.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1605.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%mul.1607.clone.1, %add.839.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.838.clone.1 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} add(%param_0.1271, %mul.1605.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.182 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} multiply(%add.838.clone.1, %add.838.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1259 = f32[]{:T(128)} constant(0) + %reduce.149 = f32[]{:T(128)} reduce(%square.182, %constant.1259), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.153.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1259), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.140 = (f32[]{:T(128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.149, %add.838.clone.1, %add.842.clone.1, %add.843.clone.1, %reduce.153.clone.1) } %region_7.10 (reduce_sum.93: f32[], reduce_sum.94: f32[]) -> f32[] { @@ -680,1321 +632,1369 @@ 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} - %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} +%fused_computation.316 (param_0.1292: bf16[128,128256], param_1.1427: bf16[128]) -> f32[128] { + %param_0.1292 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.868 = f32[128,128256]{1,0:T(8,128)} convert(%param_0.1292), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.496 = f32[1,128,128256]{2,1,0:T(8,128)} bitcast(%convert_element_type.868), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1427 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %sub.100 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1427), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.71 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%bitcast.496, %sub.100), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.56 = f32[1,128,128256]{2,1,0:T(8,128)} exponential(%sub.71), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %constant.1281 = f32[]{:T(128)} constant(0) + ROOT %reduce.150 = f32[128]{0:T(128)S(1)} reduce(%exp.56, %constant.1281), dimensions={0,2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%region_44.49 (reduce_sum.262: f32[], reduce_sum.266: f32[]) -> f32[] { + %reduce_sum.262 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.266 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + 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":[]}} } -%region_31.36 (reduce_sum.184: f32[], reduce_sum.185: f32[]) -> f32[] { +%fused_computation.322 (param_0.1291: bf16[128,128256], param_1.1425: f32[128], param_2.1107: bf16[128], param_3.748: s32[128]) -> f32[128] { + %param_3.748 = s32[128]{0:T(128)S(1)} parameter(3) + %eq.26 = s32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.748), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.51 = s32[1,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.21 = pred[1,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.26, %iota.51), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %param_0.1291 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.881 = f32[128,128256]{1,0:T(8,128)} convert(%param_0.1291), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.508 = f32[1,128,128256]{2,1,0:T(8,128)} bitcast(%convert_element_type.881), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_2.1107 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %sub.99 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1107), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.92 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%bitcast.508, %sub.99), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1425 = f32[128]{0:T(128)S(1)} parameter(1) + %sub.97 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1425), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.91 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%sub.92, %sub.97), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %constant.1279 = f32[]{:T(128)} constant(0) + %broadcast.596 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%constant.1279), dimensions={}, metadata={op_name="broadcast.104"} + %mul.1423 = f32[1,128,128256]{2,1,0:T(8,128)} select(%eq.21, %sub.91, %broadcast.596), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.151 = f32[128]{0:T(128)S(1)} reduce(%mul.1423, %constant.1279), dimensions={0,2}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%region_28.33 (reduce_sum.184: f32[], reduce_sum.185: f32[]) -> f32[] { %reduce_sum.184 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.185 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.189 = f32[]{:T(128)} add(%reduce_sum.184, %reduce_sum.185), 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":[]}} } -%region_28.33 (reduce_sum.169: f32[], reduce_sum.170: f32[]) -> f32[] { +%region_25.30 (reduce_sum.169: f32[], reduce_sum.170: f32[]) -> f32[] { %reduce_sum.169 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.170 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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} - ROOT %tuple.156 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.141, %reduce.142.clone.1) +%fused_computation.330 (param_0.1285: f32[4,1024,8,128], param_1.1421: f32[4,1024,8,128]) -> (f32[], f32[]) { + %param_0.1285 = f32[4,1024,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.527 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1285), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.185 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.527, %bitcast.527), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1273 = f32[]{:T(128)} constant(0) + %reduce.154 = f32[]{:T(128)} reduce(%square.185, %constant.1273), dimensions={0,1,2,3}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1421 = f32[4,1024,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) + %bitcast.531.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1421), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.188.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.531.clone.1, %bitcast.531.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.155.clone.1 = f32[]{:T(128)} reduce(%square.188.clone.1, %constant.1273), dimensions={0,1,2,3}, to_apply=%region_25.30, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.155 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.154, %reduce.155.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.333 (param_0.933: f32[1024,4,8,128]) -> bf16[4,1024,8,128] { + %param_0.933 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.239 = bf16[1024,4,8,128]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.933), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\'].value"} + ROOT %bitcast.532 = bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.239), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} } -%region_58.63 (reduce_sum.324: f32[], reduce_sum.325: f32[]) -> f32[] { +%region_55.60 (reduce_sum.324: f32[], reduce_sum.325: f32[]) -> f32[] { %reduce_sum.324 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.325 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.322 = f32[]{:T(128)} add(%reduce_sum.324, %reduce_sum.325), 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":[]}} } -%region_44.49 (reduce_sum.247: f32[], reduce_sum.248: f32[]) -> f32[] { +%region_41.46 (reduce_sum.247: f32[], reduce_sum.248: f32[]) -> f32[] { %reduce_sum.247 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.248 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %broadcast.557.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.920.clone.1), dimensions={}, metadata={op_name="broadcast.52"} - %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} - %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) -} - -%region_55.60 (reduce_sum.304: f32[], reduce_sum.308: f32[]) -> f32[] { +%fused_computation.334 (param_0.1269: f32[1024,4,8,128], param_1.1411: f32[], param_2.1095: f32[], param_3.737: f32[], param_4.454: f32[1024,4,8,128], param_5.400: f32[], param_6.255: f32[4,1024,8,128], param_7.171: pred[], param_8.100: f32[1024,4,8,128]) -> (f32[], f32[1024,4,8,128], f32[1024,4,8,128], f32[1024,4,8,128], f32[]) { + %param_0.1269 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %param_3.737 = f32[]{:T(128)S(6)} parameter(3) + %mul.1590.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.737), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.171 = pred[]{:T(512)S(6)} parameter(7) + %select_n.262.clone.1 = pred[1024,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.171), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.255 = f32[4,1024,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.615.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.255), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.400 = f32[]{:T(128)} parameter(5) + %div.779.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.400), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.778.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.615.clone.1, %div.779.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.261.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.615.clone.1, %div.778.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1210.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.667.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1210.clone.1), dimensions={}, metadata={op_name="broadcast.68"} + %mul.1594.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.667.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.100 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %constant.1214.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.666.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1214.clone.1), dimensions={}, metadata={op_name="broadcast.67"} + %mul.1593.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.100, %broadcast.666.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.831.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1594.clone.1, %mul.1593.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1095 = f32[]{:T(128)S(6)} parameter(2) + %div.775.clone.1 = f32[1024,4,8,128]{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.65.clone.1 = f32[1024,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.1213.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.665.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1213.clone.1), dimensions={}, metadata={op_name="broadcast.58"} + %mul.1592.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.665.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.454 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(4) + %constant.1212.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.664.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1212.clone.1), dimensions={}, metadata={op_name="broadcast.57"} + %mul.1591.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.454, %broadcast.664.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.830.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%mul.1592.clone.1, %mul.1591.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1411 = f32[]{:T(128)S(6)} parameter(1) + %div.774.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1411), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.773.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} divide(%add.830.clone.1, %div.774.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.63.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.773.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1211.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.662.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1211.clone.1), dimensions={}, metadata={op_name="broadcast.54"} + %add.829.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.662.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.282.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.775.clone.1, %add.829.clone.1), metadata={op_name="multiply.39"} + %div.772.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} divide(%add.831.clone.1, %multiply.282.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1589.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1269, %broadcast.667.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.828.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} add(%div.772.clone.1, %mul.1589.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1588.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1590.clone.1, %add.828.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.827.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1269, %mul.1588.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.189 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.827.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1257 = f32[]{:T(128)} constant(0) + %reduce.156 = f32[]{:T(128)} reduce(%square.189, %constant.1257), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.158.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1257), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.141 = (f32[]{:T(128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.156, %add.827.clone.1, %add.830.clone.1, %add.831.clone.1, %reduce.158.clone.1) +} + +%region_52.57 (reduce_sum.304: f32[], reduce_sum.308: f32[]) -> f32[] { %reduce_sum.304 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.308 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.309 = f32[]{:T(128)} add(%reduce_sum.304, %reduce_sum.308), 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":[]}} } -%region_41.46 (reduce_sum.232: f32[], reduce_sum.233: f32[]) -> f32[] { +%region_38.43 (reduce_sum.232: f32[], reduce_sum.233: f32[]) -> f32[] { %reduce_sum.232 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.233 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %broadcast.567.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.938.clone.1), dimensions={}, metadata={op_name="broadcast.52"} - %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} - %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.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.335 (param_0.1272: f32[1024,4,8,128], param_1.1414: f32[], param_2.1098: f32[], param_3.740: f32[], param_4.457: f32[1024,4,8,128], param_5.403: f32[], param_6.258: f32[4,1024,8,128], param_7.174: pred[], param_8.103: f32[1024,4,8,128]) -> (f32[], f32[1024,4,8,128], f32[1024,4,8,128], f32[1024,4,8,128], f32[]) { + %param_0.1272 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %param_3.740 = f32[]{:T(128)S(6)} parameter(3) + %mul.1617.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.740), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.174 = pred[]{:T(512)S(6)} parameter(7) + %select_n.274.clone.1 = pred[1024,4,8,128]{3,2,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.258 = f32[4,1024,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.621.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.258), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.403 = f32[]{:T(128)} parameter(5) + %div.803.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.403), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.802.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.621.clone.1, %div.803.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.273.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.274.clone.1, %bitcast.621.clone.1, %div.802.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1228.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.677.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1228.clone.1), dimensions={}, metadata={op_name="broadcast.68"} + %mul.1621.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.273.clone.1, %broadcast.677.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.103 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %constant.1232.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.676.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1232.clone.1), dimensions={}, metadata={op_name="broadcast.67"} + %mul.1620.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.103, %broadcast.676.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.848.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1621.clone.1, %mul.1620.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1098 = f32[]{:T(128)S(6)} parameter(2) + %div.799.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1098), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.68.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,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.1231.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.675.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1231.clone.1), dimensions={}, metadata={op_name="broadcast.58"} + %mul.1619.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.68.clone.1, %broadcast.675.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.457 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(4) + %constant.1230.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.674.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1230.clone.1), dimensions={}, metadata={op_name="broadcast.57"} + %mul.1618.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.457, %broadcast.674.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.847.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%mul.1619.clone.1, %mul.1618.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1414 = f32[]{:T(128)S(6)} parameter(1) + %div.798.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1414), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.797.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} divide(%add.847.clone.1, %div.798.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.66.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.797.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1229.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.672.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1229.clone.1), dimensions={}, metadata={op_name="broadcast.54"} + %add.846.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.66.clone.1, %broadcast.672.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.285.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.799.clone.1, %add.846.clone.1), metadata={op_name="multiply.36"} + %div.796.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} divide(%add.848.clone.1, %multiply.285.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1616.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1272, %broadcast.677.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.845.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} add(%div.796.clone.1, %mul.1616.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1615.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1617.clone.1, %add.845.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.844.clone.1 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1272, %mul.1615.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.190 = f32[1024,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.844.clone.1, %add.844.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1260 = f32[]{:T(128)} constant(0) + %reduce.157 = f32[]{:T(128)} reduce(%square.190, %constant.1260), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.159.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1260), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.142 = (f32[]{:T(128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.157, %add.844.clone.1, %add.847.clone.1, %add.848.clone.1, %reduce.159.clone.1) +} + +%fused_computation.351 (param_0.1002: bf16[1,128,4096], param_1.1053: f32[128], param_2.700: f32[128], param_3.393: bf16[1,128,4096], param_4.231: bf16[4096]) -> bf16[1,128,4096] { + %param_3.393 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.231 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %dot_general.398 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.231), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %dot_general.384 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.393, %dot_general.398), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0} + %convert_element_type.899 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.384), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_2.700 = f32[128]{0:T(128)S(1)} parameter(2) + %mul.1467 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.700), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1459 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.899, %mul.1467), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1002 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.910 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1002), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1053 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1466 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1053), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1465 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.910, %mul.1466), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %add_any.153 = f32[1,128,4096]{2,1,0:T(8,128)} add(%mul.1459, %mul.1465), metadata={op_name="jit(train_step)/transpose(jvp())/add_any" stack_frame_id=0} + ROOT %convert_element_type.897 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.153), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} } %region_5.8 (reduce_sum.87: f32[], reduce_sum.88: f32[]) -> f32[] { - %reduce_sum.87 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - %reduce_sum.88 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - 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":[]}} + %reduce_sum.87 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.88 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.92 = f32[]{:T(128)} add(%reduce_sum.87, %reduce_sum.88), 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.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} +%fused_computation.352 (param_0.1293: bf16[1,128,4096]) -> f32[128] { + %param_0.1293 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.901 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1293), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %square.193 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.901, %convert_element_type.901), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.1282 = f32[]{:T(128)} constant(0) + ROOT %reduce.160 = f32[128]{0:T(128)S(1)} reduce(%square.193, %constant.1282), dimensions={0,2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_10.13 (reduce_sum.102: f32[], reduce_sum.106: f32[]) -> f32[] { - %reduce_sum.102 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - %reduce_sum.106 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - 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} - %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} + %reduce_sum.102 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum"} + %reduce_sum.106 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum"} + ROOT %reduce_sum.107 = f32[]{:T(128)} add(%reduce_sum.102, %reduce_sum.106), metadata={op_name="jit(train_step)/transpose(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.1288: bf16[1,128,4096], param_1.1423: bf16[1,128,4096], param_2.1105: bf16[4096]) -> f32[128] { + %param_0.1288 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.908 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1288), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1423 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1105 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.397 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1105), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %dot_general.383 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1423, %dot_general.397), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0} + %convert_element_type.907 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.383), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %mul.1463 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.908, %convert_element_type.907), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.1276 = f32[]{:T(128)} constant(0) + ROOT %reduce.161 = f32[128]{0:T(128)S(1)} reduce(%mul.1463, %constant.1276), dimensions={0,2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0} +} + +%fused_computation.356 (param_0.1017: bf16[4096], param_1.1071: f32[128], param_2.708: bf16[1,128,4096]) -> bf16[128,4096] { + %param_0.1017 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.396 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1017), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %param_2.708 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.916 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_2.708), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1071 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1474 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1071), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1473 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.916, %mul.1474), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.915 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1473), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %dot_general.388 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.396, %convert_element_type.915), metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + ROOT %bitcast.547 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%dot_general.388) +} + +%fused_computation.326.clone (param_0.1178: bf16[128,128256], param_1.1300: f32[128], param_2.923: f32[128], param_3.557: bf16[128], param_4.307: s32[128], param_5.257: f32[128]) -> bf16[128,128256] { + %param_5.257 = f32[128]{0:T(128)S(1)} parameter(5) + %mul.1539 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.257), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_1.1300 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1538 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1300), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1178 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.947 = f32[128,128256]{1,0:T(8,128)} convert(%param_0.1178), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.589 = f32[1,128,128256]{2,1,0:T(8,128)} bitcast(%convert_element_type.947), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_3.557 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(3) + %sub.103 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.557), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.102 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%bitcast.589, %sub.103), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.67 = f32[1,128,128256]{2,1,0:T(8,128)} exponential(%sub.102), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %mul.1537 = f32[1,128,128256]{2,1,0:T(8,128)} multiply(%mul.1538, %exp.67), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.923 = f32[128]{0:T(128)S(1)} parameter(2) + %div.693 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.923), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.692 = f32[1,128,128256]{2,1,0:T(8,128)} divide(%mul.1537, %div.693), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_4.307 = s32[128]{0:T(128)S(1)} parameter(4) + %eq.29 = s32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.307), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.57 = s32[1,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.28 = pred[1,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.29, %iota.57), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %convert_element_type.946 = f32[1,128,128256]{2,1,0:T(8,128)} convert(%eq.28), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.101 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%div.692, %convert_element_type.946), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1536 = f32[1,128,128256]{2,1,0:T(8,128)} multiply(%mul.1539, %sub.101), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %convert_element_type.945 = bf16[1,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1536), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + ROOT %bitcast.588 = bf16[128,128256]{1,0:T(8,128)(2,1)} bitcast(%convert_element_type.945) +} + +%fused_computation.355 (param_0.1179: bf16[4096], param_1.1301: bf16[1,128,4096], param_2.924: f32[128], param_3.558: bf16[128,128256], param_4.308: f32[128], param_5.258: f32[128], param_6.118: bf16[128], param_7.74: s32[128], param_8: f32[128]) -> bf16[4096,128256] { + %param_0.1179 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %param_2.924 = f32[128]{0:T(128)S(1)} parameter(2) + %param_1.1301 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.298 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_0.1179, %param_2.924, %param_1.1301), kind=kLoop, calls=%fused_computation.356 + %param_3.558 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.308 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.258 = f32[128]{0:T(128)S(1)} parameter(5) + %param_6.118 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %param_7.74 = s32[128]{0:T(128)S(1)} parameter(7) + %param_8 = f32[128]{0:T(128)S(1)} parameter(8) + %fusion.319 = bf16[128,128256]{1,0:T(8,128)(2,1)} fusion(%param_3.558, %param_4.308, %param_5.258, %param_6.118, %param_7.74, /*index=5*/%param_8), kind=kLoop, calls=%fused_computation.326.clone + ROOT %convolution.118 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convolution(%fusion.298, %fusion.319), dim_labels=fb_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} } %region_8.11 (dot_general.182: bf16[], dot_general.183: bf16[]) -> bf16[] { - %dot_general.182 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general"} - %dot_general.183 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general"} - 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.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.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.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.86.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} - %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} - %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) -} - -%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.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.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.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} -} - -%region_27.32 (reduce_sum.163: f32[], reduce_sum.164: f32[]) -> f32[] { + %dot_general.182 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + %dot_general.183 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + ROOT %add.165 = bf16[]{:T(256)} add(%dot_general.182, %dot_general.183), metadata={op_name="add.51"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.326.clone.1.clone.clone (param_0.1263: bf16[128,128256], param_1.1407: f32[128], param_2.1084: f32[128], param_3.732: bf16[128], param_4.450: s32[128], param_5.396: f32[128]) -> bf16[128,128256] { + %param_5.396 = f32[128]{0:T(128)S(1)} parameter(5) + %mul.1667 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.396), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_1.1407 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1666 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1407), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1263 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.966 = f32[128,128256]{1,0:T(8,128)} convert(%param_0.1263), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.653 = f32[1,128,128256]{2,1,0:T(8,128)} bitcast(%convert_element_type.966), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_3.732 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(3) + %sub.114 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.732), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.113 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%bitcast.653, %sub.114), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.72 = f32[1,128,128256]{2,1,0:T(8,128)} exponential(%sub.113), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %mul.1665 = f32[1,128,128256]{2,1,0:T(8,128)} multiply(%mul.1666, %exp.72), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1084 = f32[128]{0:T(128)S(1)} parameter(2) + %div.833 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1084), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.832 = f32[1,128,128256]{2,1,0:T(8,128)} divide(%mul.1665, %div.833), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_4.450 = s32[128]{0:T(128)S(1)} parameter(4) + %eq.35 = s32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.450), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.60 = s32[1,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.34 = pred[1,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.35, %iota.60), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %convert_element_type.965 = f32[1,128,128256]{2,1,0:T(8,128)} convert(%eq.34), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.112 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%div.832, %convert_element_type.965), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1664 = f32[1,128,128256]{2,1,0:T(8,128)} multiply(%mul.1667, %sub.112), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %convert_element_type.964 = bf16[1,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1664), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + ROOT %bitcast.652 = bf16[128,128256]{1,0:T(8,128)(2,1)} bitcast(%convert_element_type.964) +} + +%bitcast_fusion.4 (bitcast_input.4: bf16[4096,128256]) -> bf16[4096,128256] { + %bitcast_input.4 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.660 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) +} + +%fused_computation.359 (param_0.1262: f32[128], param_1.1406: bf16[1,128,4096], param_2.1085: bf16[4096,128256], param_3.733: bf16[128,128256], param_4.451: f32[128], param_5.397: f32[128], param_6.252: bf16[128], param_7.168: s32[128], param_8.97: f32[128]) -> (bf16[4096], bf16[1,128,4096]) { + %param_3.733 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.451 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.397 = f32[128]{0:T(128)S(1)} parameter(5) + %param_6.252 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %param_7.168 = s32[128]{0:T(128)S(1)} parameter(7) + %param_8.97 = f32[128]{0:T(128)S(1)} parameter(8) + %fusion.320.clone.1 = bf16[128,128256]{1,0:T(8,128)(2,1)} fusion(%param_3.733, %param_4.451, %param_5.397, %param_6.252, %param_7.168, /*index=5*/%param_8.97), kind=kLoop, calls=%fused_computation.326.clone.1.clone.clone + %param_2.1085 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(2) + %fusion.332 = bf16[4096,128256]{1,0:T(8,128)(2,1)} fusion(%param_2.1085), kind=kLoop, calls=%bitcast_fusion.4 + %convolution.117.clone.1 = bf16[128,4096]{1,0:T(8,128)(2,1)} convolution(%fusion.320.clone.1, %fusion.332), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} + %bitcast.545.clone.1 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.117.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} + %param_1.1406 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.920 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_1.1406), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_0.1262 = f32[128]{0:T(128)S(1)} parameter(0) + %mul.1478 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1262), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1477 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.920, %mul.1478), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.919 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1477), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %multiply.271 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.545.clone.1, %convert_element_type.919), metadata={op_name="multiply.197"} + %constant.1136 = bf16[]{:T(256)} constant(0) + %reduce.162 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%multiply.271, %constant.1136), dimensions={0,1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0} + ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)S(1)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.162, %bitcast.545.clone.1) +} + +%fused_computation.369 (param_0.1070: f32[128], param_1.1137: f32[64]) -> (bf16[1,128,1,64], bf16[1,128,1,64]) { + %param_0.1070 = f32[128]{0:T(128)S(1)} parameter(0) + %div.631 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_0.1070), dimensions={1}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1137 = f32[64]{0:T(128)S(1)} parameter(1) + %div.639 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_1.1137), dimensions={3}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.626 = f32[1,128,1,64]{3,1,2,0:T(8,128)} divide(%div.631, %div.639), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sin.38 = f32[1,128,1,64]{3,1,2,0:T(8,128)} sine(%div.626), metadata={op_name="jit(train_step)/sin" stack_frame_id=0} + %convert_element_type.928 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %cos.41.clone.1 = f32[1,128,1,64]{3,1,2,0:T(8,128)} cosine(%div.626), metadata={op_name="jit(train_step)/cos" stack_frame_id=0} + %convert_element_type.927.clone.1 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + ROOT %tuple.148 = (bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.928, %convert_element_type.927.clone.1) +} + +%region_24.29 (reduce_sum.163: f32[], reduce_sum.164: f32[]) -> f32[] { %reduce_sum.163 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.164 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.168 = f32[]{:T(128)} add(%reduce_sum.163, %reduce_sum.164), 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":[]}} } -%region_26.31 (reduce_sum.157: f32[], reduce_sum.161: f32[]) -> f32[] { +%fused_computation.371 (param_0.1279: f32[4,4096]) -> f32[] { + %param_0.1279 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.553 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1279), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.196 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.553, %bitcast.553), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1267 = f32[]{:T(128)} constant(0) + ROOT %reduce.163 = f32[]{:T(128)} reduce(%square.196, %constant.1267), dimensions={0,1}, to_apply=%region_24.29, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_23.28 (reduce_sum.157: f32[], reduce_sum.161: f32[]) -> f32[] { %reduce_sum.157 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.161 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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} - ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.150, %reduce.151.clone.1) +%fused_computation.373 (param_0.1280: f32[4,4096]) -> f32[] { + %param_0.1280 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.557 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1280), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.199 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.557, %bitcast.557), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1268 = f32[]{:T(128)} constant(0) + ROOT %reduce.164 = f32[]{:T(128)} reduce(%square.199, %constant.1268), dimensions={0,1}, to_apply=%region_23.28, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%region_54.59 (reduce_sum.301: f32[], reduce_sum.302: f32[]) -> f32[] { +%region_51.56 (reduce_sum.301: f32[], reduce_sum.302: f32[]) -> f32[] { %reduce_sum.301 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.302 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.303 = f32[]{:T(128)} add(%reduce_sum.301, %reduce_sum.302), 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":[]}} } -%region_40.45 (reduce_sum.226: f32[], reduce_sum.227: f32[]) -> f32[] { +%region_37.42 (reduce_sum.226: f32[], reduce_sum.227: f32[]) -> f32[] { %reduce_sum.226 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.227 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %broadcast.573.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.944.clone.1), dimensions={}, metadata={op_name="broadcast.53"} - %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} - %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.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) -} - -%region_53.58 (reduce_sum.295: f32[], reduce_sum.296: f32[]) -> f32[] { +%fused_computation.374 (param_0.1273: f32[4096,4], param_1.1415: f32[], param_2.1099: f32[], param_3.741: f32[], param_4.458: f32[4096,4], param_5.404: f32[], param_6.259: f32[4,4096], param_7.175: pred[], param_8.104: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1273 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.741 = f32[]{:T(128)S(6)} parameter(3) + %mul.1546.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.741), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.175 = pred[]{:T(512)S(6)} parameter(7) + %select_n.242.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.175), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.259 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.611.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.259), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.404 = f32[]{:T(128)} parameter(5) + %div.739.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.404), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.738.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.611.clone.1, %div.739.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.241.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.242.clone.1, %bitcast.611.clone.1, %div.738.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1180.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.649.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1180.clone.1), dimensions={}, metadata={op_name="broadcast.70"} + %mul.1550.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.241.clone.1, %broadcast.649.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.104 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %constant.1184.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.648.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1184.clone.1), dimensions={}, metadata={op_name="broadcast.69"} + %mul.1549.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.104, %broadcast.648.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.803.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} 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.735.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1099), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.60.clone.1 = f32[4096,4]{0,1:T(4,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.1183.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.647.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1183.clone.1), dimensions={}, metadata={op_name="broadcast.60"} + %mul.1548.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.60.clone.1, %broadcast.647.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.458 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %constant.1182.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.646.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1182.clone.1), dimensions={}, metadata={op_name="broadcast.59"} + %mul.1547.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.458, %broadcast.646.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.802.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1548.clone.1, %mul.1547.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1415 = f32[]{:T(128)S(6)} parameter(1) + %div.734.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.733.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.802.clone.1, %div.734.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.58.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.733.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1181.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.644.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1181.clone.1), dimensions={}, metadata={op_name="broadcast.55"} + %add.801.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.58.clone.1, %broadcast.644.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.277.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.735.clone.1, %add.801.clone.1), metadata={op_name="multiply.35"} + %div.732.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.803.clone.1, %multiply.277.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1545.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1273, %broadcast.649.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.800.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.732.clone.1, %mul.1545.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1544.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1546.clone.1, %add.800.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.799.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1273, %mul.1544.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.799.clone.1, %add.799.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1261 = f32[]{:T(128)} constant(0) + %reduce.165 = f32[]{:T(128)} reduce(%square.200, %constant.1261), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.167.clone.1 = f32[]{:T(128)} reduce(%integer_pow.60.clone.1, %constant.1261), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.143 = (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.165, %add.799.clone.1, %add.802.clone.1, %add.803.clone.1, %reduce.167.clone.1) +} + +%region_50.55 (reduce_sum.295: f32[], reduce_sum.296: f32[]) -> f32[] { %reduce_sum.295 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.296 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.297 = f32[]{:T(128)} add(%reduce_sum.295, %reduce_sum.296), 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":[]}} } -%region_39.44 (reduce_sum.220: f32[], reduce_sum.224: f32[]) -> f32[] { +%region_36.41 (reduce_sum.220: f32[], reduce_sum.224: f32[]) -> f32[] { %reduce_sum.220 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.224 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %broadcast.579.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.950.clone.1), dimensions={}, metadata={op_name="broadcast.53"} - %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} - %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.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) +%fused_computation.375 (param_0.1274: f32[4096,4], param_1.1416: f32[], param_2.1100: f32[], param_3.742: f32[], param_4.459: f32[4096,4], param_5.405: f32[], param_6.260: f32[4,4096], param_7.176: pred[], param_8.105: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { + %param_0.1274 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.742 = f32[]{:T(128)S(6)} parameter(3) + %mul.1553.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.742), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.176 = pred[]{:T(512)S(6)} parameter(7) + %select_n.246.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.176), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.260 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.613.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.260), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.405 = f32[]{:T(128)} parameter(5) + %div.747.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.405), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.746.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.613.clone.1, %div.747.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.245.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.246.clone.1, %bitcast.613.clone.1, %div.746.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1186.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.655.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1186.clone.1), dimensions={}, metadata={op_name="broadcast.70"} + %mul.1557.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.245.clone.1, %broadcast.655.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.105 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) + %constant.1190.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.654.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1190.clone.1), dimensions={}, metadata={op_name="broadcast.69"} + %mul.1556.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.105, %broadcast.654.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.1557.clone.1, %mul.1556.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) + %div.743.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.61.clone.1 = f32[4096,4]{0,1:T(4,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.1189.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.653.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1189.clone.1), dimensions={}, metadata={op_name="broadcast.60"} + %mul.1555.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.61.clone.1, %broadcast.653.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.459 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) + %constant.1188.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.652.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1188.clone.1), dimensions={}, metadata={op_name="broadcast.59"} + %mul.1554.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.459, %broadcast.652.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.807.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1555.clone.1, %mul.1554.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1416 = f32[]{:T(128)S(6)} parameter(1) + %div.742.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.741.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.807.clone.1, %div.742.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.59.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.741.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1187.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.650.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.1187.clone.1), dimensions={}, metadata={op_name="broadcast.55"} + %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.59.clone.1, %broadcast.650.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.278.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.743.clone.1, %add.806.clone.1), metadata={op_name="multiply.34"} + %div.740.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.808.clone.1, %multiply.278.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1552.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1274, %broadcast.655.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)} add(%div.740.clone.1, %mul.1552.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1551.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1553.clone.1, %add.805.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.804.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1274, %mul.1551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.201 = f32[4096,4]{0,1:T(4,128)} multiply(%add.804.clone.1, %add.804.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1262 = f32[]{:T(128)} constant(0) + %reduce.166 = f32[]{:T(128)} reduce(%square.201, %constant.1262), dimensions={0,1}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.168.clone.1 = f32[]{:T(128)} reduce(%integer_pow.61.clone.1, %constant.1262), dimensions={0,1}, to_apply=%region_36.41, 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.166, %add.804.clone.1, %add.807.clone.1, %add.808.clone.1, %reduce.168.clone.1) +} + +%fused_computation.384 (param_0.1065: bf16[1,128,1,64]) -> bf16[1,128,1,128] { + %param_0.1065 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1128 = bf16[]{:T(256)} constant(-inf) + %pad.38 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1065, %constant.1128), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %pad.37 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1065, %constant.1128), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + ROOT %maximum.38 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.38, %pad.37), metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} +} + +%fused_computation.385 (param_0.1067: bf16[1,128,1,64]) -> bf16[1,128,1,128] { + %param_0.1067 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1127 = bf16[]{:T(256)} constant(-inf) + %pad.40 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1067, %constant.1127), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %pad.39 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1067, %constant.1127), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + ROOT %maximum.39 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.40, %pad.39), metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} } %region_9.12 (reduce_sum.99: f32[], reduce_sum.100: f32[]) -> f32[] { - %reduce_sum.100 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.99 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.100 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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} +%fused_computation.388 (param_0.1289: bf16[4096]) -> f32[] { + %param_0.1289 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.932 = f32[4096]{0:T(1024)} convert(%param_0.1289), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %square.204 = f32[4096]{0:T(1024)} multiply(%convert_element_type.932, %convert_element_type.932), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1277 = f32[]{:T(128)} constant(0) + ROOT %reduce.169 = f32[]{:T(128)} reduce(%square.204, %constant.1277), 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[] { +%region_46.51 (reduce_sum.274: f32[], reduce_sum.275: f32[]) -> f32[] { %reduce_sum.274 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.275 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.276 = f32[]{:T(128)} add(%reduce_sum.274, %reduce_sum.275), 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":[]}} } -%region_35.40 (reduce_sum.199: f32[], reduce_sum.203: f32[]) -> f32[] { +%region_32.37 (reduce_sum.199: f32[], reduce_sum.203: f32[]) -> f32[] { %reduce_sum.199 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.203 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - %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) - %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} - %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) - %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} - %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) - %add.834.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.974.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - %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.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] { - %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) - %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} - %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) - %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} - %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} -} - -%region_61.66 (reduce_sum.345: f32[], reduce_sum.346: f32[]) -> f32[] { +%fused_computation.389 (param_0.1278: f32[4096], param_1.1420: f32[], param_2.1104: f32[], param_3.746: f32[], param_4.463: f32[4096], param_5.409: f32[], param_6.264: bf16[4096], param_7.180: pred[], param_8.109: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { + %param_0.1278 = f32[4096]{0:T(1024)S(1)} parameter(0) + %param_3.746 = f32[]{:T(128)S(6)} parameter(3) + %mul.1560.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.746), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.180 = pred[]{:T(512)S(6)} parameter(7) + %select_n.250.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.180), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.264 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.952.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.264), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_5.409 = f32[]{:T(128)} parameter(5) + %div.755.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.409), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.754.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.952.clone.1, %div.755.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.249.clone.1 = f32[4096]{0:T(1024)} select(%select_n.250.clone.1, %convert_element_type.952.clone.1, %div.754.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1192.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.657.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.1192.clone.1), dimensions={}, metadata={op_name="broadcast.74"} + %mul.1566.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.249.clone.1, %broadcast.657.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.109 = f32[4096]{0:T(1024)S(1)} parameter(8) + %constant.1196.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.1567.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.1196.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1565.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.109, %mul.1567.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.814.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1566.clone.1, %mul.1565.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1104 = f32[]{:T(128)S(6)} parameter(2) + %div.751.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1104), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.62.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.249.clone.1, %select_n.249.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.1195.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.1564.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.1195.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1562.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.62.clone.1, %mul.1564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.463 = f32[4096]{0:T(1024)S(1)} parameter(4) + %constant.1194.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.1563.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.1194.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.1561.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.463, %mul.1563.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.813.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1562.clone.1, %mul.1561.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1420 = f32[]{:T(128)S(6)} parameter(1) + %div.750.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1420), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.749.clone.1 = f32[4096]{0:T(1024)} divide(%add.813.clone.1, %div.750.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.60.clone.1 = f32[4096]{0:T(1024)} sqrt(%div.749.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1193.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.812.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.1193.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.811.clone.1 = f32[4096]{0:T(1024)} add(%sqrt.60.clone.1, %add.812.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.279.clone.1 = f32[4096]{0:T(1024)} multiply(%div.751.clone.1, %add.811.clone.1), metadata={op_name="multiply.30"} + %div.748.clone.1 = f32[4096]{0:T(1024)} divide(%add.814.clone.1, %multiply.279.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1559.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1278, %broadcast.657.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.810.clone.1 = f32[4096]{0:T(1024)} add(%div.748.clone.1, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1558.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1560.clone.1, %add.810.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.809.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1278, %mul.1558.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.205 = f32[4096]{0:T(1024)} multiply(%add.809.clone.1, %add.809.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1266 = f32[]{:T(128)} constant(0) + %reduce.170 = f32[]{:T(128)} reduce(%square.205, %constant.1266), dimensions={0}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.171.clone.1 = f32[]{:T(128)} reduce(%integer_pow.62.clone.1, %constant.1266), dimensions={0}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.147 = (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.170, %add.809.clone.1, %add.813.clone.1, %add.814.clone.1, %reduce.171.clone.1) +} + +%fused_computation.394 (param_0.1095: s32[512]) -> s32[1024] { + %constant.1063 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.624 = s32[1024]{0:T(1024)} broadcast(%constant.1063), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %param_0.1095 = s32[512]{0:T(512)S(1)} parameter(0) + %constant.1064 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.1095, %constant.1064), padding=0_512, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %constant.1062 = s32[] constant(128255), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.623 = s32[1024]{0:T(1024)} broadcast(%constant.1062), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %clamp.3 = s32[1024]{0:T(1024)S(1)} clamp(%broadcast.624, %pad.41, %broadcast.623), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} +} + +%fused_computation.397 (param_0.1158: f32[128]) -> f32[128] { + %param_0.1158 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1132 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.632 = f32[128]{0:T(128)} broadcast(%constant.1132), dimensions={}, metadata={op_name="broadcast.283"} + %div.667 = f32[128]{0:T(128)} multiply(%param_0.1158, %broadcast.632), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.1130 = f32[]{:T(128)} constant(1e-05) + %broadcast.630 = f32[128]{0:T(128)} broadcast(%constant.1130), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.791 = f32[128]{0:T(128)} add(%div.667, %broadcast.630), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %bitcast.568 = f32[1,128]{1,0:T(1,128)} bitcast(%add.791), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %rsqrt.101 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.568), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} + ROOT %bitcast.558 = f32[128]{0:T(128)} bitcast(%rsqrt.101), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} +} + +%region_0.1 (reduce_sum.67: s32[], reduce_sum.71: s32[]) -> s32[] { + %reduce_sum.67 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.71 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + 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.402 (param_0.1166: s32[1,128]) -> s32[] { + %param_0.1166 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.1153 = s32[]{:T(128)} constant(0) + %broadcast.643 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1153), dimensions={}, metadata={op_name="broadcast.85"} + %ne.10 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1166, %broadcast.643), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %convert_element_type.941 = s32[1,128]{1,0:T(1,128)} convert(%ne.10), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + ROOT %reduce.172 = s32[]{:T(128)} reduce(%convert_element_type.941, %constant.1153), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%region_58.63 (reduce_sum.345: f32[], reduce_sum.346: f32[]) -> f32[] { %reduce_sum.345 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} %reduce_sum.346 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} ROOT %reduce_sum.330 = f32[]{:T(128)} add(%reduce_sum.345, %reduce_sum.346), 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":[]}} } -%region_48.53 (reduce_sum.268: f32[], reduce_sum.269: f32[]) -> f32[] { +%region_45.50 (reduce_sum.268: f32[], reduce_sum.269: f32[]) -> f32[] { %reduce_sum.268 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} %reduce_sum.269 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} 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) - %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} - %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} - %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} - 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) - %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} - %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} - %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} +%fused_computation.405 (param_0.1290: s32[1,128], param_1.1424: bf16[128], param_2.1106: f32[128], param_3.747: f32[128], param_4.464: f32[]) -> (f32[], f32[], f32[128]) { + %param_0.1290 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.1150 = s32[]{:T(128)} constant(0) + %broadcast.641 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1150), dimensions={}, metadata={op_name="broadcast.85"} + %ne.16 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1290, %broadcast.641), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %param_2.1106 = f32[128]{0:T(128)S(1)} parameter(2) + %log.18 = f32[128]{0:T(128)} log(%param_2.1106), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_1.1424 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %reduce_max.16 = f32[128]{0:T(128)} convert(%param_1.1424), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %add.794 = f32[128]{0:T(128)} add(%log.18, %reduce_max.16), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %square.208 = f32[128]{0:T(128)} multiply(%add.794, %add.794), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.1278 = f32[]{:T(128)} constant(0) + %broadcast.626 = f32[128]{0:T(128)} broadcast(%constant.1278), dimensions={}, metadata={op_name="broadcast.94"} + %mul.1531 = f32[128]{0:T(128)} multiply(%square.208, %broadcast.626), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %bitcast.561 = f32[1,128]{1,0:T(1,128)} bitcast(%mul.1531), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %broadcast.637 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1278), dimensions={}, metadata={op_name="broadcast.94"} + %mul.1517 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %bitcast.561, %broadcast.637), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.174 = f32[]{:T(128)} reduce(%mul.1517, %constant.1278), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_3.747 = f32[128]{0:T(128)S(1)} parameter(3) + %neg.118.clone.1 = f32[128]{0:T(128)} negate(%param_3.747), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %bitcast.565.clone.1 = f32[1,128]{1,0:T(1,128)} bitcast(%neg.118.clone.1), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.784.clone.1 = f32[1,128]{1,0:T(1,128)} add(%bitcast.565.clone.1, %bitcast.561), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1515.clone.1 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %add.784.clone.1, %broadcast.637), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.173.clone.1 = f32[]{:T(128)} reduce(%mul.1515.clone.1, %constant.1278), dimensions={0,1}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_4.464 = f32[]{:T(128)S(6)} parameter(4) + %broadcast_in_dim.295.clone.1 = f32[1,128]{1,0:T(1,128)} broadcast(%param_4.464), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} + %mul.1513.clone.1 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %broadcast_in_dim.295.clone.1, %broadcast.637), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %bitcast.559.clone.1 = f32[128]{0:T(128)S(1)} bitcast(%mul.1513.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %tuple.149 = (f32[]{:T(128)}, f32[]{:T(128)}, f32[128]{0:T(128)S(1)}) tuple(%reduce.174, %reduce.173.clone.1, %bitcast.559.clone.1) +} + +%fused_computation.408 (param_0.1130: s32[1,128]) -> s32[1,1,128] { + %param_0.1130 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.1154 = s32[]{:T(128)} constant(0) + %broadcast.642 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1154), dimensions={}, metadata={op_name="broadcast.85"} + %lt.32 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1130, %broadcast.642), direction=LT, metadata={op_name="jit(train_step)/jvp()/lt" stack_frame_id=0} + %constant.1137 = s32[]{:T(128)} constant(128256) + %add.792 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1137), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.781 = s32[1,128]{1,0:T(1,128)} add(%param_0.1130, %add.792), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %select_n.178 = s32[1,128]{1,0:T(1,128)} select(%lt.32, %add.781, %param_0.1130), metadata={op_name="jit(train_step)/jvp()/select_n" stack_frame_id=0} + ROOT %bitcast.560 = s32[1,1,128]{2,1,0:T(1,128)S(1)} bitcast(%select_n.178) +} + +%fused_computation.414 (param_0.1160: f32[128], param_1.1244: f32[128]) -> f32[128] { + %param_0.1160 = f32[128]{0:T(128)S(1)} parameter(0) + %bitcast.575 = f32[1,128]{1,0:T(1,128)} bitcast(%param_0.1160), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0} + %param_1.1244 = f32[128]{0:T(128)S(1)} parameter(1) + %constant.1131 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.631 = f32[128]{0:T(128)} broadcast(%constant.1131), dimensions={}, metadata={op_name="broadcast.283"} + %div.665 = f32[128]{0:T(128)} multiply(%param_1.1244, %broadcast.631), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.1129 = f32[]{:T(128)} constant(1e-05) + %broadcast.629 = f32[128]{0:T(128)} broadcast(%constant.1129), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.789 = f32[128]{0:T(128)} add(%div.665, %broadcast.629), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %bitcast.574 = f32[1,128]{1,0:T(1,128)} bitcast(%add.789), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %rsqrt.107 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.574), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} + %div.663 = f32[1,128]{1,0:T(1,128)} divide(%rsqrt.107, %bitcast.574), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.1126 = f32[]{:T(128)} constant(-0.5) + %mul.1535 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1126), dimensions={}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1529 = f32[1,128]{1,0:T(1,128)} multiply(%div.663, %mul.1535), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1528 = f32[1,128]{1,0:T(1,128)} multiply(%bitcast.575, %mul.1529), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.1125 = f32[]{:T(128)} constant(0.00048828125) + %mul.1534 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1125), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1524 = f32[1,128]{1,0:T(1,128)} multiply(%mul.1528, %mul.1534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %bitcast.569 = f32[128]{0:T(128)S(1)} bitcast(%mul.1524), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +} + +%fused_computation.415 () -> f32[64] { + %constant.1135 = f32[]{:T(128)} constant(500000) + %broadcast.635 = f32[64]{0:T(128)} broadcast(%constant.1135), dimensions={}, metadata={op_name="broadcast.274"} + %iota.56 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/iota" stack_frame_id=0} + %constant.1134 = s32[]{:T(128)} constant(2) + %broadcast.634 = s32[64]{0:T(128)} broadcast(%constant.1134), dimensions={}, metadata={op_name="broadcast.275"} + %mul.1525 = s32[64]{0:T(128)} multiply(%iota.56, %broadcast.634), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %convert_element_type.942 = f32[64]{0:T(128)} convert(%mul.1525), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %constant.1133 = f32[]{:T(128)} constant(0.0078125) + %broadcast.633 = f32[64]{0:T(128)} broadcast(%constant.1133), dimensions={}, metadata={op_name="broadcast.276"} + %div.661 = f32[64]{0:T(128)} multiply(%convert_element_type.942, %broadcast.633), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + ROOT %pow.36 = f32[64]{0:T(128)} power(%broadcast.635, %div.661), metadata={op_name="jit(train_step)/pow" stack_frame_id=0} +} + +%fused_computation.416 (param_0.1146: s32[1,128]) -> (f32[1,128,1,1], f32[128]) { + %param_0.1146 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %convert_element_type.943 = f32[1,128]{1,0:T(1,128)} convert(%param_0.1146), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %bitcast.570 = f32[1,128,1,1]{1,3,2,0:T(1,128)} bitcast(%convert_element_type.943), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.571.clone.1 = f32[128]{0:T(128)S(1)} bitcast(%convert_element_type.943), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + ROOT %tuple.151 = (f32[1,128,1,1]{1,3,2,0:T(1,128)}, f32[128]{0:T(128)S(1)}) tuple(%bitcast.570, %bitcast.571.clone.1) +} + +%region_33.38.clone (reduce_sum.520: f32[], reduce_sum.521: f32[]) -> f32[] { + %reduce_sum.520 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.521 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.522 = f32[]{:T(128)} add(%reduce_sum.520, %reduce_sum.521), 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":[]}} +} + +%region_20.25.clone (reduce_sum.422: f32[], reduce_sum.423: f32[]) -> f32[] { + %reduce_sum.422 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.423 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.427 = f32[]{:T(128)} add(%reduce_sum.422, %reduce_sum.423), 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.420 (param_0.1264: f32[4096,4]) -> bf16[4,4096] { + %param_0.1264 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.654 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1264), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.654) +} + +%fused_computation.421 (param_0.1265: f32[4096,4]) -> bf16[4,4096] { + %param_0.1265 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.655 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1265), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.110 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.655) +} + +%fused_computation.422 (param_0.1256: f32[128], param_1.1426: bf16[128]) -> (f32[128], f32[128]) { + %param_0.1256 = f32[128]{0:T(128)S(1)} parameter(0) + %log.23 = f32[128]{0:T(128)S(1)} log(%param_0.1256), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_1.1426 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %reduce_max.18.clone.1 = f32[128]{0:T(128)} convert(%param_1.1426), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %add.796.clone.1 = f32[128]{0:T(128)} add(%log.23, %reduce_max.18.clone.1), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %constant.1280 = f32[]{:T(128)} constant(0) + %broadcast.628.clone.1 = f32[128]{0:T(128)} broadcast(%constant.1280), dimensions={}, metadata={op_name="broadcast.94"} + %mul.1523.clone.1 = f32[128]{0:T(128)} multiply(%add.796.clone.1, %broadcast.628.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.1155.clone.1 = f32[]{:T(128)} constant(1) + %broadcast.625.clone.1 = f32[128]{0:T(128)} broadcast(%constant.1155.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.785.clone.1 = f32[128]{0:T(128)S(1)} add(%mul.1523.clone.1, %broadcast.625.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + ROOT %tuple.150 = (f32[128]{0:T(128)S(1)}, f32[128]{0:T(128)S(1)}) tuple(%log.23, %add.785.clone.1) } -%region_0.1 (reduce_sum.67: s32[], reduce_sum.71: s32[]) -> s32[] { - %reduce_sum.67 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.71 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - 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"]}]}} +%region_6.9 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { + %reduce_max.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_max"} + %reduce_max.8 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_max"} + 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.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.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.358.clone.clone (param_0.1260: bf16[4096], param_1.1403: f32[128], param_2.1079: bf16[1,128,4096]) -> bf16[128,4096] { + %param_0.1260 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.406 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1260), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %param_2.1079 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.960 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_2.1079), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1403 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1659 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1403), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1658 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.960, %mul.1659), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.959 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1658), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %dot_general.405 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.406, %convert_element_type.959), metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + ROOT %bitcast.649 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%dot_general.405) } -%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.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} +%bitcast_fusion.5 (bitcast_input.5: bf16[4096,128256]) -> bf16[4096,128256] { + %bitcast_input.5 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.661 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.5) } -%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.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.423 (param_0.1266: bf16[4096,128256], param_1.1408: bf16[4096], param_2.1092: f32[128], param_3.734: bf16[1,128,4096]) -> (bf16[128], bf16[128,128256]) { + %param_1.1408 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1092 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.734 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %fusion.300.clone.1 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1408, %param_2.1092, %param_3.734), kind=kLoop, calls=%fused_computation.358.clone.clone + %param_0.1266 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(0) + %fusion.333 = bf16[4096,128256]{1,0:T(8,128)(2,1)} fusion(%param_0.1266), kind=kLoop, calls=%bitcast_fusion.5 + %convolution.119.clone.1 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} convolution(%fusion.300.clone.1, %fusion.333), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/dot_general" stack_frame_id=0} + %constant.1254 = bf16[]{:T(256)} constant(-inf) + %reduce.175 = bf16[128]{0:T(256)(128)(2,1)S(1)} reduce(%convolution.119.clone.1, %constant.1254), dimensions={1}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + ROOT %tuple.152 = (bf16[128]{0:T(256)(128)(2,1)S(1)}, bf16[128,128256]{1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.175, %convolution.119.clone.1) } -%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.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.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} +%convert_element_type.525.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { + %lhs.1 = bf16[] parameter(0) + %rhs.1 = bf16[] parameter(1) + ROOT %add.660 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%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.238.clone.clone (param_0.1428: bf16[4,4096], param_1.1512: s32[]) -> bf16[4096] { + %param_0.1428 = bf16[4,4096]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %param_1.1512 = s32[]{:T(128)S(6)} parameter(1) + %constant.1392 = s32[]{:T(128)} constant(0) + %dynamic_slice.281 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1428, %param_1.1512, %constant.1392), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.1393 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %reduce.189 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.281, %constant.1393), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} } -%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.99.clone.clone (param_0.1442: bf16[4,1024,32,128], param_1.1521: s32[]) -> bf16[1,1024,32,128] { + %param_0.1442 = bf16[4,1024,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1521 = s32[]{:T(128)S(6)} parameter(1) + %constant.1404 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.172 = bf16[1,1024,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1442, %param_1.1521, %constant.1404, %constant.1404, %constant.1404), dynamic_slice_sizes={1,1024,32,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.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_12.14 (reduce_sum.108: f32[], reduce_sum.109: f32[]) -> f32[] { + %reduce_sum.108 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.109 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.113 = f32[]{:T(128)} add(%reduce_sum.108, %reduce_sum.109), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.153.clone.clone (param_0.1425: bf16[4,1,128,4096], param_1.1510: s32[]) -> f32[128] { + %param_0.1425 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.1510 = s32[]{:T(128)S(6)} parameter(1) + %constant.1386 = s32[]{:T(128)} constant(0) + %dynamic-slice.167 = bf16[1,1,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1425, %param_1.1510, %constant.1386, %constant.1386, %constant.1386), dynamic_slice_sizes={1,1,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.784 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.167), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1015 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.784), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.215 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1015, %convert_element_type.1015), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.1387 = f32[]{:T(128)} constant(0) + ROOT %reduce.187 = f32[128]{0:T(128)S(1)} reduce(%square.215, %constant.1387), dimensions={0,2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} +} + +%fused_computation.243.clone.clone (param_0.1426: f32[128]) -> f32[128] { + %param_0.1426 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1389 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.709 = f32[128]{0:T(128)} broadcast(%constant.1389), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.852 = f32[128]{0:T(128)} multiply(%param_0.1426, %broadcast.709), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1388 = f32[]{:T(128)} constant(1e-05) + %broadcast.708 = f32[128]{0:T(128)} broadcast(%constant.1388), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.886 = f32[128]{0:T(128)} add(%div.852, %broadcast.708), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.786 = f32[1,128]{1,0:T(1,128)} bitcast(%add.886), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.114 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.786), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.785 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.114), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.165.clone.clone (param_0.1446: bf16[1,128,32,128]) -> (bf16[1,128,32,64], bf16[1,128,32,64]) { + %param_0.1446 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.75 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1446), slice={[0:1], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %neg.132 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%slice.75), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0} + %slice.76 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1446), slice={[0:1], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.186 = (bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.132, %slice.76) +} + +%fused_computation.252.clone.clone () -> f32[64] { + %constant.1398 = f32[]{:T(128)} constant(500000) + %closed_call.64 = f32[64]{0:T(128)} broadcast(%constant.1398), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %iota.65 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/iota" stack_frame_id=0} + %constant.1397 = s32[]{:T(128)} constant(2) + %closed_call.52 = s32[64]{0:T(128)} broadcast(%constant.1397), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.1753 = s32[64]{0:T(128)} multiply(%iota.65, %closed_call.52), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1018 = f32[64]{0:T(128)} convert(%mul.1753), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %constant.1396 = f32[]{:T(128)} constant(0.0078125) + %closed_call.51 = f32[64]{0:T(128)} broadcast(%constant.1396), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.853 = f32[64]{0:T(128)} multiply(%convert_element_type.1018, %closed_call.51), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.64, %div.853), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/pow" stack_frame_id=0} +} + +%fused_computation.234.clone.clone (param_0.1433: f32[128], param_1.1516: f32[64]) -> (bf16[1,128,1,64], bf16[1,128,1,64]) { + %param_0.1433 = f32[128]{0:T(128)S(1)} parameter(0) + %div.855 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_0.1433), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %param_1.1516 = f32[64]{0:T(128)S(1)} parameter(1) + %div.856 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_1.1516), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %div.854 = f32[1,128,1,64]{3,1,2,0:T(8,128)} divide(%div.855, %div.856), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %cos.43 = f32[1,128,1,64]{3,1,2,0:T(8,128)} cosine(%div.854), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/cos" stack_frame_id=0} + %convert_element_type.1019 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %sin.35.clone.3 = f32[1,128,1,64]{3,1,2,0:T(8,128)} sine(%div.854), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/sin" stack_frame_id=0} + %convert_element_type.768.clone.3 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %tuple.184 = (bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1019, %convert_element_type.768.clone.3) +} + +%fused_computation.237.clone.clone (param_0.1435: bf16[1,128,1,64]) -> bf16[128,128] { + %param_0.1435 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1400 = bf16[]{:T(256)} constant(-inf) + %pad.61 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1435, %constant.1400), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %pad.60 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1435, %constant.1400), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %maximum.49 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.61, %pad.60), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + ROOT %bitcast.792 = bf16[128,128]{1,0:T(8,128)(2,1)} bitcast(%maximum.49), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} +} + +%fused_computation.236.clone.clone (param_0.1434: bf16[1,128,1,64]) -> bf16[128,128] { + %param_0.1434 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1399 = bf16[]{:T(256)} constant(-inf) + %pad.59 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1434, %constant.1399), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %pad.58 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1434, %constant.1399), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %maximum.48 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.59, %pad.58), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + ROOT %bitcast.791 = bf16[128,128]{1,0:T(8,128)(2,1)} bitcast(%maximum.48), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} +} + +%fused_computation.168.clone.clone (param_0.1447: bf16[1,128,32,64], param_1.1524: bf16[1,128,32,64], param_2.1168: bf16[1,128,32,128], param_3.790: bf16[128,128], param_4.488: bf16[128,128]) -> bf16[32,128,128] { + %param_2.1168 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.488 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1765 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.488), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1763 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1168, %mul.1765), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %param_1.1524 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1406 = bf16[]{:T(256)} constant(-inf) + %pad.65 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1524, %constant.1406), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %param_0.1447 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.64 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1447, %constant.1406), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %maximum.51 = bf16[1,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())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %param_3.790 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1764 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.790), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1762 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.51, %mul.1764), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %add.888 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1763, %mul.1762), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + ROOT %bitcast.802 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.888), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(jit(_splash_attention))/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.117.clone.clone (param_0.1436: bf16[4,1024,8,128], param_1.1517: s32[]) -> bf16[1,1024,8,128] { + %param_0.1436 = bf16[4,1024,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1517 = s32[]{:T(128)S(6)} parameter(1) + %constant.1401 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.170 = bf16[1,1024,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1436, %param_1.1517, %constant.1401, %constant.1401, %constant.1401), dynamic_slice_sizes={1,1024,8,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.221.clone.clone (param_0.1440: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1440 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.73 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1440), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %neg.131 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%slice.73), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0} + %slice.74 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1440), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.185 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.131, %slice.74) +} + +%fused_computation.224.clone.clone (param_0.1441: bf16[1,128,8,64], param_1.1520: bf16[1,128,8,64], param_2.1165: bf16[1,128,8,128], param_3.787: bf16[128,128], param_4.486: bf16[128,128]) -> bf16[8,128,128] { + %param_2.1165 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.486 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1759 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.486), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1757 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1165, %mul.1759), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %param_1.1520 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1403 = bf16[]{:T(256)} constant(-inf) + %pad.63 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1520, %constant.1403), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %param_0.1441 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.62 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1441, %constant.1403), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %maximum.50 = bf16[1,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())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %param_3.787 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1758 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.787), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1756 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.50, %mul.1758), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %add.887 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1757, %mul.1756), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + ROOT %bitcast.797 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)} bitcast(%add.887), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(jit(_splash_attention))/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.116.clone.clone (param_0.1429: bf16[4,1024,8,128], param_1.1513: s32[]) -> bf16[1,1024,8,128] { + %param_0.1429 = bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1513 = s32[]{:T(128)S(6)} parameter(1) + %constant.1394 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.168 = bf16[1,1024,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1429, %param_1.1513, %constant.1394, %constant.1394, %constant.1394), dynamic_slice_sizes={1,1024,8,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.253.clone.clone (param_0.1476: f32[1,32,128,128]) -> (f32[32,128], f32[1,32,128,1]) { + %param_0.1476 = f32[1,32,128,128]{2,1,3,0:T(8,128)} parameter(0) + %slice.77 = f32[1,32,128,1]{2,1,3,0:T(8,128)} slice(%param_0.1476), slice={[0:1], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.827 = f32[32,128]{1,0:T(8,128)} bitcast(%slice.77), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/broadcast_in_dim" stack_frame_id=0} + ROOT %tuple.191 = (f32[32,128]{1,0:T(8,128)}, f32[1,32,128,1]{2,1,3,0:T(8,128)}) tuple(%bitcast.827, %slice.77) +} + +%fused_computation.98.clone.clone (param_0.1448: bf16[4,32,128,1024], param_1.1525: s32[]) -> bf16[1,32,128,1024] { + %param_0.1448 = bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1525 = s32[]{:T(128)S(6)} parameter(1) + %constant.1407 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.174 = bf16[1,32,128,1024]{3,2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1448, %param_1.1525, %constant.1407, %constant.1407, %constant.1407), dynamic_slice_sizes={1,32,128,1024}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} } -%region_6.9 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { - %reduce_max.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_max"} - %reduce_max.8 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_max"} - 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":[]}} +%convert_element_type.523.reduce_sub_computation (lhs: bf16[], rhs: bf16[]) -> bf16[] { + %lhs = bf16[] parameter(0) + %rhs = bf16[] parameter(1) + ROOT %add.659 = bf16[] add(%lhs, %rhs), 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.239.clone.clone (param_0.1427: bf16[4,4096], param_1.1511: s32[]) -> bf16[4096] { + %param_0.1427 = bf16[4,4096]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %param_1.1511 = s32[]{:T(128)S(6)} parameter(1) + %constant.1390 = s32[]{:T(128)} constant(0) + %dynamic_slice.280 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1427, %param_1.1511, %constant.1390), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.1391 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %reduce.188 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%dynamic_slice.280, %constant.1391), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/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.84.clone.clone (param_0.1453: bf16[4,1024,14336], param_1.1527: s32[]) -> bf16[1,1024,14336] { + %param_0.1453 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1527 = s32[]{:T(128)S(6)} parameter(1) + %constant.1411 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.175 = bf16[1,1024,14336]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1453, %param_1.1527, %constant.1411, %constant.1411), dynamic_slice_sizes={1,1024,14336}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} } -%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.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.85.clone.clone (param_0.1457: bf16[4,14336,1024], param_1.1530: s32[]) -> bf16[1,14336,1024] { + %param_0.1457 = bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1530 = s32[]{:T(128)S(6)} parameter(1) + %constant.1412 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.176 = bf16[1,14336,1024]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1457, %param_1.1530, %constant.1412, %constant.1412), dynamic_slice_sizes={1,14336,1024}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} } -%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.76.clone.clone.clone.clone (param_0.1459: bf16[1,14336,4096]) -> bf16[14336,4096] { + %param_0.1459 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.811 = bf16[14336,4096]{1,0:T(8,128)(2,1)} bitcast(%param_0.1459), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} } -%convert_element_type.525.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { - %lhs.1 = bf16[] parameter(0) - %rhs.1 = bf16[] parameter(1) - 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.180.clone.clone (param_0.1458: bf16[1,128,4096]) -> bf16[128,4096] { + %param_0.1458 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.810 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%param_0.1458) } -%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.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.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} +%fused_computation.109.clone.clone (param_0.1460: bf16[1,14336,4096], param_1.1531: bf16[1,128,4096]) -> bf16[1,128,14336] { + %param_0.1460 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.429 = bf16[14336,4096]{1,0:T(8,128)(2,1)} fusion(%param_0.1460), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %param_1.1531 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.430 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1531), kind=kLoop, calls=%fused_computation.180.clone.clone + %convolution.146 = bf16[14336,128]{0,1:T(8,128)(2,1)} convolution(%fusion.429, %fusion.430), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %bitcast.812 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.146), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} } -%region_12.14 (reduce_sum.108: f32[], reduce_sum.109: f32[]) -> f32[] { - %reduce_sum.108 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.109 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - 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) - %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} - %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) - %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} - %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) - %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} -} - -%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.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.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) -} - -%fused_computation.145.clone.clone () -> f32[64] { - %constant.1124 = f32[]{:T(128)} constant(500000) - %closed_call.84 = f32[64]{0:T(128)} broadcast(%constant.1124), 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.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} - %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} - 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} - %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} - %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) -} - -%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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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[] { - %reduce_sum.120 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.121 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - 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) - %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} -} - -%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.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.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.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) +%fused_computation.86.clone.clone (param_0.1461: bf16[4,1024,14336], param_1.1532: s32[]) -> bf16[1,1024,14336] { + %param_0.1461 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1532 = s32[]{:T(128)S(6)} parameter(1) + %constant.1413 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.177 = bf16[1,1024,14336]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1461, %param_1.1532, %constant.1413, %constant.1413), dynamic_slice_sizes={1,1024,14336}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} } -%convert_element_type.523.reduce_sub_computation (lhs: bf16[], rhs: bf16[]) -> bf16[] { - %lhs = bf16[] parameter(0) - %rhs = bf16[] parameter(1) - 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) - %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":[]} - %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} -} - -%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.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.565 = 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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.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.241.clone.clone (param_0.1452: f32[128]) -> f32[128] { + %param_0.1452 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1410 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.711 = f32[128]{0:T(128)} broadcast(%constant.1410), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.857 = f32[128]{0:T(128)} multiply(%param_0.1452, %broadcast.711), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1409 = f32[]{:T(128)} constant(1e-05) + %broadcast.710 = f32[128]{0:T(128)} broadcast(%constant.1409), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.889 = f32[128]{0:T(128)} add(%div.857, %broadcast.710), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.806 = f32[1,128]{1,0:T(1,128)} bitcast(%add.889), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.115 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.806), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.805 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.115), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.74.clone.clone (param_0.1465: bf16[1,4096,14336]) -> bf16[4096,14336] { + %param_0.1465 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.816 = bf16[4096,14336]{1,0:T(8,128)(2,1)} bitcast(%param_0.1465), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.125.clone.clone (param_0.1466: bf16[1,128,14336], param_1.1535: bf16[1,128,14336]) -> bf16[128,14336] { + %param_1.1535 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1414 = bf16[]{:T(256)} constant(1) + %jit_silu_.42 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1414), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)" stack_frame_id=0} + %neg.133 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1535), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/neg" stack_frame_id=0} + %exp.79 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.133), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/exp" stack_frame_id=0} + %add.890 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.79, %jit_silu_.42), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/add" stack_frame_id=0} + %div.858 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.42, %add.890), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/div" stack_frame_id=0} + %mul.1771 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1535, %div.858), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/mul" stack_frame_id=0} + %param_0.1466 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %mul.1770 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1771, %param_0.1466), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + ROOT %bitcast.817 = bf16[128,14336]{1,0:T(8,128)(2,1)} bitcast(%mul.1770) +} + +%fused_computation.73.clone.clone (param_0.1467: bf16[1,4096,14336], param_1.1536: bf16[1,128,14336], param_2.1174: bf16[1,128,14336]) -> bf16[128,4096] { + %param_1.1536 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1174 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.434 = bf16[128,14336]{1,0:T(8,128)(2,1)} fusion(%param_1.1536, %param_2.1174), kind=kLoop, calls=%fused_computation.125.clone.clone + %param_0.1467 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.433 = bf16[4096,14336]{1,0:T(8,128)(2,1)} fusion(%param_0.1467), kind=kLoop, calls=%fused_computation.74.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %convolution.148 = bf16[128,4096]{1,0:T(8,128)(2,1)S(1)} convolution(%fusion.434, %fusion.433), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} } %region_14.17 (reduce_sum.126: f32[], reduce_sum.127: f32[]) -> f32[] { - %reduce_sum.126 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum"} - %reduce_sum.127 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum"} - 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) - %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} -} - -%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.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.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.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.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.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.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} - %rsqrt.99 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.864), 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} - %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} - %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} -} - -%region_20.24 (dot_general.187: bf16[], dot_general.188: bf16[]) -> bf16[] { - %dot_general.187 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general"} - %dot_general.188 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general"} - 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"} - %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} - %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} - %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} - 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) + %reduce_sum.126 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + %reduce_sum.127 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + ROOT %reduce_sum.128 = f32[]{:T(128)} add(%reduce_sum.126, %reduce_sum.127), metadata={op_name="checkpoint/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.78.clone.clone.clone.clone.clone.clone.clone.clone (param_0.1468: bf16[1,4096,14336]) -> bf16[4096,14336] { + %param_0.1468 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.818 = bf16[4096,14336]{1,0:T(8,128)(2,1)} bitcast(%param_0.1468), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.120.clone.1.clone.clone.clone.clone (param_0.1469: bf16[1,128,14336], param_1.1537: bf16[1,128,14336], param_2.1175: bf16[1,128,14336]) -> bf16[128,14336] { + %param_1.1537 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1175 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.1776 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1537, %param_2.1175), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.1415 = bf16[]{:T(256)} constant(1) + %jit_silu_.43 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1415), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)" stack_frame_id=0} + %param_0.1469 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %neg.134 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1469), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/neg" stack_frame_id=0} + %exp.80 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.134), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/exp" stack_frame_id=0} + %add.891 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.80, %jit_silu_.43), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/add" stack_frame_id=0} + %div.859 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.43, %add.891), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/div" stack_frame_id=0} + %mul.1775 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1776, %div.859), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0} + %mul.1774 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1469, %mul.1776), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0} + %sub.118 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} subtract(%jit_silu_.43, %div.859), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/sub" stack_frame_id=0} + %mul.1773 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.859, %sub.118), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/mul" stack_frame_id=0} + %mul.1772 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1774, %mul.1773), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0} + %add_any.160 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1775, %mul.1772), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/add_any" stack_frame_id=0} + ROOT %bitcast.819 = bf16[128,14336]{1,0:T(8,128)(2,1)} bitcast(%add_any.160) +} + +%fused_computation.158.clone.clone (param_0.1470: bf16[1,128,4096], param_1.1538: bf16[4096], param_2.1176: bf16[128,4096], param_3.794: bf16[1,4096,14336], param_4.489: bf16[1,128,14336], param_5.416: bf16[1,128,14336], param_6.268: bf16[1,128,14336]) -> (f32[128], bf16[1,128,4096]) { + %param_0.1470 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1030 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1470), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_2.1176 = bf16[128,4096]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_4.489 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %param_5.416 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) + %param_6.268 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) + %fusion.271.clone.3 = bf16[128,14336]{1,0:T(8,128)(2,1)} fusion(%param_4.489, %param_5.416, %param_6.268), kind=kLoop, calls=%fused_computation.120.clone.1.clone.clone.clone.clone + %param_3.794 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) + %fusion.158.clone.3 = bf16[4096,14336]{1,0:T(8,128)(2,1)} fusion(%param_3.794), kind=kLoop, calls=%fused_computation.78.clone.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.97.clone.3 = bf16[128,4096]{1,0:T(8,128)(2,1)} convolution(%fusion.271.clone.3, %fusion.158.clone.3), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %add_any.122.clone.3 = bf16[128,4096]{1,0:T(8,128)(2,1)} add(%param_2.1176, %convolution.97.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %bitcast.323.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%add_any.122.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %param_1.1538 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %dot_general.453 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1538), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.452 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.323.clone.3, %dot_general.453), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1029 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.452), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %mul.1777 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1030, %convert_element_type.1029), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.1416 = f32[]{:T(128)} constant(0) + %reduce.191 = f32[128]{0:T(128)S(1)} reduce(%mul.1777, %constant.1416), dimensions={0,2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum" stack_frame_id=0} + ROOT %tuple.188 = (f32[128]{0:T(128)S(1)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.191, %bitcast.323.clone.3) +} + +%fused_computation.251.clone.clone (param_0.1471: f32[128], param_1.1539: f32[128]) -> f32[128] { + %param_0.1471 = f32[128]{0:T(128)S(1)} parameter(0) + %bitcast.822 = f32[1,128]{1,0:T(1,128)} bitcast(%param_0.1471), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum" stack_frame_id=0} + %param_1.1539 = f32[128]{0:T(128)S(1)} parameter(1) + %constant.1420 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.713 = f32[128]{0:T(128)} broadcast(%constant.1420), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.861 = f32[128]{0:T(128)} multiply(%param_1.1539, %broadcast.713), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1419 = f32[]{:T(128)} constant(1e-05) + %broadcast.712 = f32[128]{0:T(128)} broadcast(%constant.1419), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.892 = f32[128]{0:T(128)} add(%div.861, %broadcast.712), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.821 = f32[1,128]{1,0:T(1,128)} bitcast(%add.892), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.116 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.821), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + %div.860 = f32[1,128]{1,0:T(1,128)} divide(%rsqrt.116, %bitcast.821), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1418 = f32[]{:T(128)} constant(-0.5) + %closed_call.65 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1418), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.1780 = f32[1,128]{1,0:T(1,128)} multiply(%div.860, %closed_call.65), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1779 = f32[1,128]{1,0:T(1,128)} multiply(%bitcast.822, %mul.1780), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.1417 = f32[]{:T(128)} constant(0.00048828125) + %mul.1781 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1417), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.1778 = f32[1,128]{1,0:T(1,128)} multiply(%mul.1779, %mul.1781), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + ROOT %bitcast.820 = f32[128]{0:T(128)S(1)} bitcast(%mul.1778), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%region_17.21 (dot_general.187: bf16[], dot_general.188: bf16[]) -> bf16[] { + %dot_general.187 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + %dot_general.188 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + ROOT %add.167 = bf16[]{:T(256)} add(%dot_general.187, %dot_general.188), metadata={op_name="add.44"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.196.clone.clone (param_0.1472: bf16[1,128,4096], param_1.1540: f32[128], param_2.1177: bf16[1,128,4096], param_3.795: bf16[1,128,4096], param_4.490: f32[128], param_5.417: bf16[4096]) -> (bf16[4096], bf16[1,128,4096]) { + %param_0.1472 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_2.1177 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1032 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_2.1177), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1540 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1783 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1540), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1782 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1032, %mul.1783), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1031 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1782), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %multiply.298 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1472, %convert_element_type.1031), metadata={op_name="multiply.195"} + %constant.1421 = bf16[]{:T(256)} constant(0) + %reduce.192 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.298, %constant.1421), dimensions={0,1}, to_apply=%region_17.21, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.795 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_5.417 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %dot_general.378.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.417), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.294.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1472, %dot_general.378.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.766.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.294.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %mul.1210.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.766.clone.3, %mul.1783), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %param_4.490 = f32[128]{0:T(128)S(1)} parameter(4) + %mul.1223.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.490), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.1209.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1032, %mul.1223.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %add_any.142.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} add(%mul.1210.clone.3, %mul.1209.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %convert_element_type.764.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.142.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %add_any.140.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.795, %convert_element_type.764.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + ROOT %tuple.189 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.192, %add_any.140.clone.3) } %region_15.18 (dot_general.184: f32[], dot_general.185: f32[]) -> f32[] { %dot_general.184 = f32[]{:T(128)} parameter(0), metadata={op_name="vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general"} %dot_general.185 = f32[]{:T(128)} parameter(1), metadata={op_name="vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general"} - 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":[]}} -} + ROOT %add.166 = 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.88.clone.clone.clone.clone (param_0.1473: bf16[1,32,128,4096]) -> bf16[32,128,4096] { + %param_0.1473 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.823 = bf16[32,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1473), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.178.clone.clone.clone.clone (param_0.1474: bf16[1,128,4096]) -> bf16[128,4096,1] { + %param_0.1474 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.824 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1474) +} + +%fused_computation.159.clone.clone (param_0.1475: bf16[32,128,128], param_1.1541: bf16[1,32,128,4096], param_2.1178: bf16[1,128,4096]) -> (f32[32,128], bf16[128,32,128]) { + %param_0.1475 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.826 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1475), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(jit(_splash_attention))/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/broadcast_in_dim" stack_frame_id=0} + %convert.125 = f32[1,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.826), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/convert" stack_frame_id=0} + %param_2.1178 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.193.clone.3 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1178), kind=kLoop, calls=%fused_computation.178.clone.clone.clone.clone + %param_1.1541 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.107.clone.3 = bf16[32,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1541), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.70.clone.3 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)S(1)} convolution(%fusion.193.clone.3, %fusion.107.clone.3), window={size=32 pad=31_31 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %bitcast.825 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%convolution.70.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/transpose" stack_frame_id=0} + %convert.124 = f32[1,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.825), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/convert.1" stack_frame_id=0} + %multiply.299 = f32[1,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.125, %convert.124), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/multiply" stack_frame_id=0} + %constant.1422 = f32[]{:T(128)} constant(0) + %dot_general.454 = f32[32,128]{1,0:T(8,128)S(1)} reduce(%multiply.299, %constant.1422), dimensions={0,3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} + ROOT %tuple.190 = (f32[32,128]{1,0:T(8,128)S(1)}, bf16[128,32,128]{2,0,1:T(8,128)(2,1)S(1)}) tuple(%dot_general.454, %convolution.70.clone.3) +} + +%fused_computation.176.clone.clone (param_0.1488: bf16[32,128,128], param_1.1549: bf16[128,128]) -> bf16[1,128,32,128] { + %param_0.1488 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.842 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1488), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/splash_mha_dq_segmented_no_residuals/splash_mha_dq_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %param_1.1549 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(1) + %broadcast.716 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_1.1549), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1790 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} multiply(%bitcast.842, %broadcast.716), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + ROOT %bitcast.841 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%mul.1790), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%fused_computation.418.clone.clone (param_0.1489: bf16[1,128,32,128]) -> (bf16[1,128,32,64], bf16[1,128,32,64]) { + %param_0.1489 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.79 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1489), slice={[0:1], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %slice.50.clone.3 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1489), slice={[0:1], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %neg.103.clone.3 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%slice.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/neg" stack_frame_id=0} + ROOT %tuple.195 = (bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%slice.79, %neg.103.clone.3) +} + +%fused_computation.170.clone.clone (param_0.1490: bf16[1,128,32,64], param_1.1550: bf16[1,128,32,64], param_2.1185: bf16[32,128,128], param_3.802: bf16[128,128]) -> bf16[128,32,128] { + %param_1.1550 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1429 = bf16[]{:T(256)} constant(-inf) + %pad.69 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1550, %constant.1429), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %param_0.1490 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.68 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1490, %constant.1429), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %maximum.53 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.69, %pad.68), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %param_2.1185 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.845 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_2.1185), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/splash_mha_dq_segmented_no_residuals/splash_mha_dq_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %param_3.802 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %broadcast.717 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_3.802), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1791 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} multiply(%bitcast.845, %broadcast.717), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %bitcast.844 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%mul.1791), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %add_any.162 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%maximum.53, %bitcast.844), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + ROOT %bitcast.843 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)S(1)} bitcast(%add_any.162) +} + +%fused_computation.92.clone.clone.clone.clone (param_0.1501: bf16[1,4096,32,128]) -> bf16[4096,32,128] { + %param_0.1501 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.856 = bf16[4096,32,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1501), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%bitcast_fusion.1.clone.clone (bitcast_input.13: bf16[128,32,128]) -> bf16[128,32,128] { + %bitcast_input.13 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.857 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} bitcast(%bitcast_input.13) +} + +%fused_computation.141.clone.clone (param_0.1502: bf16[128,32,128], param_1.1557: bf16[1,4096,32,128]) -> bf16[128,4096] { + %param_0.1502 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.446 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1502), kind=kLoop, calls=%bitcast_fusion.1.clone.clone + %param_1.1557 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.445 = bf16[4096,32,128]{0,2,1:T(8,128)(2,1)} fusion(%param_1.1557), kind=kLoop, calls=%fused_computation.92.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.154 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.446, %fusion.445), window={size=32}, dim_labels=b0f_o0i->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %bitcast.858 = bf16[128,4096]{1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.154), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.231.clone.clone (param_0.1480: bf16[8,128,128], param_1.1544: bf16[128,128]) -> bf16[1,128,8,128] { + %param_0.1480 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.832 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1480), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/splash_mha_dkv_segmented_no_residuals/splash_mha_dkv_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %param_1.1544 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(1) + %broadcast.714 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_1.1544), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1786 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} multiply(%bitcast.832, %broadcast.714), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + ROOT %bitcast.831 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%mul.1786), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%fused_computation.419.clone.clone (param_0.1481: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1481 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.78 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1481), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %slice.52.clone.3 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1481), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %neg.105.clone.3 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%slice.52.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/neg" stack_frame_id=0} + ROOT %tuple.192 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%slice.78, %neg.105.clone.3) +} + +%fused_computation.226.clone.clone (param_0.1482: bf16[1,128,8,64], param_1.1545: bf16[1,128,8,64], param_2.1181: bf16[8,128,128], param_3.798: bf16[128,128]) -> bf16[128,8,128] { + %param_1.1545 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1424 = bf16[]{:T(256)} constant(-inf) + %pad.67 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1545, %constant.1424), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %param_0.1482 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.66 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1482, %constant.1424), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %maximum.52 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.67, %pad.66), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %param_2.1181 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.835 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_2.1181), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/splash_mha_dkv_segmented_no_residuals/splash_mha_dkv_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %param_3.798 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %broadcast.715 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_3.798), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1787 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} multiply(%bitcast.835, %broadcast.715), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %bitcast.834 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%mul.1787), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %add_any.161 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%maximum.52, %bitcast.834), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + ROOT %bitcast.833 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)S(1)} bitcast(%add_any.161) +} + +%fused_computation.105.clone.clone.clone.clone (param_0.1499: bf16[1,4096,8,128]) -> bf16[4096,8,128] { + %param_0.1499 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.853 = bf16[4096,8,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1499), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%bitcast_fusion.clone.clone (bitcast_input.12: bf16[128,8,128]) -> bf16[128,8,128] { + %bitcast_input.12 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.854 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%bitcast_input.12) +} + +%fused_computation.139.clone.clone (param_0.1500: bf16[128,8,128], param_1.1556: bf16[1,4096,8,128]) -> bf16[128,4096] { + %param_0.1500 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.444 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1500), kind=kLoop, calls=%bitcast_fusion.clone.clone + %param_1.1556 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.443 = bf16[4096,8,128]{0,2,1:T(8,128)(2,1)} fusion(%param_1.1556), kind=kLoop, calls=%fused_computation.105.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.153 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.444, %fusion.443), window={size=8}, dim_labels=b0f_o0i->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %bitcast.855 = bf16[128,4096]{1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.153), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} -%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.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.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.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.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.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) +%region_16.20 (reduce_sum.129: f32[], reduce_sum.133: f32[]) -> f32[] { + %reduce_sum.129 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + %reduce_sum.133 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + ROOT %reduce_sum.134 = f32[]{:T(128)} add(%reduce_sum.129, %reduce_sum.133), metadata={op_name="checkpoint/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } + diff --git a/tests/utils/reference_hlo_qwen3_1.7b.txt b/tests/utils/reference_hlo_qwen3_1.7b.txt index f1ede66966..3b9fc0270a 100644 --- a/tests/utils/reference_hlo_qwen3_1.7b.txt +++ b/tests/utils/reference_hlo_qwen3_1.7b.txt @@ -1,4 +1,4 @@ -HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias) }, entry_computation_layout={(s32[]{:T(128)}, f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, /*index=5*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, /*index=10*/f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, s32[]{:T(128)}, /*index=15*/f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, /*index=20*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, /*index=25*/f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, /*index=30*/f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, /*index=35*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, /*index=40*/f32[151936,2048]{1,0:T(8,128)}, s32[]{:T(128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, /*index=45*/s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)})->(s32[]{:T(128)}, f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, /*index=5*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, /*index=10*/f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, s32[]{:T(128)}, /*index=15*/f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, /*index=20*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, /*index=25*/f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, /*index=30*/f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, /*index=35*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, /*index=40*/f32[151936,2048]{1,0:T(8,128)}, s32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=45*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=50*/f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 +HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias), {42}: (42, {}, may-alias), {43}: (43, {}, may-alias), {44}: (44, {}, may-alias), {45}: (45, {}, may-alias), {46}: (46, {}, may-alias), {47}: (47, {}, may-alias), {48}: (48, {}, may-alias), {49}: (49, {}, may-alias), {50}: (50, {}, may-alias), {51}: (51, {}, may-alias), {52}: (52, {}, may-alias), {53}: (53, {}, may-alias), {54}: (54, {}, may-alias), {55}: (55, {}, may-alias), {56}: (56, {}, may-alias), {57}: (57, {}, may-alias), {58}: (58, {}, may-alias), {59}: (59, {}, may-alias), {60}: (60, {}, may-alias), {61}: (61, {}, may-alias), {62}: (62, {}, may-alias), {63}: (63, {}, may-alias), {64}: (64, {}, may-alias), {65}: (65, {}, may-alias) }, entry_computation_layout={(f32[2048]{0:T(1024)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=5*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=10*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, /*index=15*/f32[6144,4,512]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=20*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=25*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=30*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, /*index=35*/u32[4]{0:T(128)}, f32[151936,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[2048]{0:T(1024)}, f32[512,4,6144]{2,1,0:T(4,128)}, /*index=40*/f32[512,4,6144]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=45*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=50*/f32[151936,512]{1,0:T(8,128)}, f32[2048]{0:T(1024)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, /*index=55*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, /*index=60*/f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,512]{1,0:T(8,128)}, s32[]{:T(128)}, /*index=65*/u32[]{:T(128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, /*index=70*/s32[1,128]{1,0:T(1,128)})->(f32[2048]{0:T(1024)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=5*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=10*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, /*index=15*/f32[6144,4,512]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=20*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=25*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=30*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, /*index=35*/u32[4]{0:T(128)}, f32[151936,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[2048]{0:T(1024)}, f32[512,4,6144]{2,1,0:T(4,128)}, /*index=40*/f32[512,4,6144]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=45*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=50*/f32[151936,512]{1,0:T(8,128)}, f32[2048]{0:T(1024)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, /*index=55*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, /*index=60*/f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,512]{1,0:T(8,128)}, s32[]{:T(128)}, /*index=65*/u32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=70*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=75*/s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 FileNames @@ -9,1992 +9,1992 @@ FileLocations StackFrames -%fused_computation (param_0.2: bf16[151936,2048], param_1.7: s32[1024]) -> bf16[512,2048] { - %param_0.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) +%region_0.1.clone (reduce_sum.612: s32[], reduce_sum.613: s32[]) -> s32[] { + %reduce_sum.612 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.613 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.614 = s32[]{:T(128)} add(%reduce_sum.612, %reduce_sum.613), 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"]}]}} +} + +%add.15.clone (x.31: bf16[], y.31: bf16[]) -> bf16[] { + %x.31 = bf16[]{:T(256)} parameter(0) + %y.31 = bf16[]{:T(256)} parameter(1) + ROOT %add.614 = bf16[]{:T(256)} add(%x.31, %y.31), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation (param_0.2: bf16[151936,512], param_1.7: s32[1024]) -> bf16[512,512] { + %param_0.2 = bf16[151936,512]{1,0:T(8,128)(2,1)} parameter(0) %param_1.7 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.1 = s32[1024]{0:T(1024)} custom-call(%param_1.7), 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.6 = s32[512]{0:T(512)} slice(%custom-call.1), 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.444 = s32[4,128]{1,0:T(4,128)} reshape(%slice.6), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.461 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.444), 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.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} gather(%param_0.2, %transpose.461), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,2048}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %transpose.460 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} transpose(%gather.4), 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.443 = bf16[512,2048]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.460), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} + %custom-call.1 = s32[1024]{0:T(1024)} custom-call(%param_1.7), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %slice.34 = s32[512]{0:T(512)} slice(%custom-call.1), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %reshape.793 = s32[4,128]{1,0:T(4,128)} reshape(%slice.34), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.442 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.793), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %gather.4 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} gather(%param_0.2, %transpose.442), 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()/gather" stack_frame_id=0} + %transpose.441 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%gather.4), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %reshape.792 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.441), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} } -%region_42.47.clone (scatter-add.6: bf16[], scatter-add.7: bf16[]) -> bf16[] { - %scatter-add.7 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - %scatter-add.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - ROOT %add.584 = bf16[]{:T(256)} add(%scatter-add.6, %scatter-add.7), 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":[]}} +%region_39.44.clone (scatter-add.6: bf16[], scatter-add.7: bf16[]) -> bf16[] { + %scatter-add.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add"} + %scatter-add.7 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add"} + ROOT %add.615 = bf16[]{:T(256)} add(%scatter-add.6, %scatter-add.7), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.1 (param_0.3: bf16[151936,2048], param_1.5: s32[512], param_2.4: bf16[512,2048]) -> bf16[151936,2048] { - %param_0.3 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) +%fused_computation.1 (param_0.3: bf16[151936,512], param_1.5: s32[512], param_2.4: bf16[512,512]) -> bf16[151936,512] { + %param_0.3 = bf16[151936,512]{1,0:T(8,128)(2,1)} parameter(0) %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) - %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} + %reshape.800 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.5), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %transpose.447 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.800), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %param_2.4 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %reshape.801 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} reshape(%param_2.4), metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + %transpose.448 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%reshape.801), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} + ROOT %scatter.2 = bf16[151936,512]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.447, %transpose.448), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_39.44.clone, metadata={op_name="jit(train_step)/transpose(jvp())/scatter-add" stack_frame_id=0} } -%region_71.76 (reduce_sum.464: f32[], reduce_sum.465: f32[]) -> f32[] { - %reduce_sum.465 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_68.73 (reduce_sum.464: f32[], reduce_sum.465: f32[]) -> f32[] { %reduce_sum.464 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.465 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.466 = f32[]{:T(128)} add(%reduce_sum.464, %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":[]}} } -%region_56.61 (reduce_sum.386: f32[], reduce_sum.387: f32[]) -> f32[] { - %reduce_sum.387 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_53.58 (reduce_sum.386: f32[], reduce_sum.387: f32[]) -> f32[] { %reduce_sum.386 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.387 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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} - %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} - %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) - %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} - %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) - %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} - %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) - %add.939.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1093.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - 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) -} - -%region_43.48 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { - %reduce_sum.318 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.317 (param_0.1523: f32[151936,512], param_1.1704: f32[], param_2.1310: f32[], param_3.872: f32[], param_4.510: f32[151936,512], param_5.454: f32[], param_6.316: bf16[151936,512], param_7.190: bf16[151936,2048], param_8.112: s32[], param_9.94: pred[], param_10.90: f32[151936,512]) -> (f32[], f32[151936,512], f32[151936,512], f32[151936,512], f32[]) { + %param_0.1523 = f32[151936,512]{1,0:T(8,128)} parameter(0) + %param_3.872 = f32[]{:T(128)S(6)} parameter(3) + %mul.2036.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%param_3.872), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_9.94 = pred[]{:T(512)S(6)} parameter(9) + %select_n.288.clone.1 = pred[151936,512]{1,0:T(8,128)(4,1)} broadcast(%param_9.94), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_7.190 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(7) + %constant.1408.clone.1 = s32[]{:T(128)} constant(0) + %param_8.112 = s32[]{:T(128)S(6)} parameter(8) + %dynamic-slice.149.clone.1 = bf16[151936,512]{1,0:T(8,128)(2,1)} dynamic-slice(%param_7.190, %constant.1408.clone.1, %param_8.112), dynamic_slice_sizes={151936,512}, metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294965759","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %convert_element_type.1338.clone.1 = f32[151936,512]{1,0:T(8,128)} convert(%dynamic-slice.149.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_6.316 = bf16[151936,512]{1,0:T(8,128)(2,1)} parameter(6) + %convert_element_type.1337.clone.1 = f32[151936,512]{1,0:T(8,128)} convert(%param_6.316), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %add_any.197.clone.1 = f32[151936,512]{1,0:T(8,128)} add(%convert_element_type.1338.clone.1, %convert_element_type.1337.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add_any" stack_frame_id=0} + %param_5.454 = f32[]{:T(128)} parameter(5) + %div.920.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%param_5.454), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.919.clone.1 = f32[151936,512]{1,0:T(8,128)} divide(%add_any.197.clone.1, %div.920.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.287.clone.1 = f32[151936,512]{1,0:T(8,128)} select(%select_n.288.clone.1, %add_any.197.clone.1, %div.919.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1406.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1013.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%constant.1406.clone.1), dimensions={}, metadata={op_name="broadcast.76"} + %mul.2042.clone.1 = f32[151936,512]{1,0:T(8,128)} multiply(%select_n.287.clone.1, %broadcast.1013.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_10.90 = f32[151936,512]{1,0:T(8,128)} parameter(10) + %constant.1412.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.2043.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%constant.1412.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2041.clone.1 = f32[151936,512]{1,0:T(8,128)} multiply(%param_10.90, %mul.2043.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.979.clone.1 = f32[151936,512]{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.1310 = f32[]{:T(128)S(6)} parameter(2) + %div.916.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%param_2.1310), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.70.clone.1 = f32[151936,512]{1,0:T(8,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.1410.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.2040.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%constant.1410.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2038.clone.1 = f32[151936,512]{1,0:T(8,128)} multiply(%integer_pow.70.clone.1, %mul.2040.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.510 = f32[151936,512]{1,0:T(8,128)} parameter(4) + %constant.1409.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.2039.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%constant.1409.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2037.clone.1 = f32[151936,512]{1,0:T(8,128)} multiply(%param_4.510, %mul.2039.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.978.clone.1 = f32[151936,512]{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.1704 = f32[]{:T(128)S(6)} parameter(1) + %div.915.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%param_1.1704), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.914.clone.1 = f32[151936,512]{1,0:T(8,128)} divide(%add.978.clone.1, %div.915.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.67.clone.1 = f32[151936,512]{1,0:T(8,128)} sqrt(%div.914.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1407.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.977.clone.1 = f32[151936,512]{1,0:T(8,128)} broadcast(%constant.1407.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.976.clone.1 = f32[151936,512]{1,0:T(8,128)} add(%sqrt.67.clone.1, %add.977.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.453.clone.1 = f32[151936,512]{1,0:T(8,128)} multiply(%div.916.clone.1, %add.976.clone.1), metadata={op_name="multiply.60"} + %div.913.clone.1 = f32[151936,512]{1,0:T(8,128)} divide(%add.979.clone.1, %multiply.453.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2035.clone.1 = f32[151936,512]{1,0:T(8,128)} multiply(%param_0.1523, %broadcast.1013.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.975.clone.1 = f32[151936,512]{1,0:T(8,128)} add(%div.913.clone.1, %mul.2035.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2034.clone.1 = f32[151936,512]{1,0:T(8,128)} multiply(%mul.2036.clone.1, %add.975.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.974.clone.1 = f32[151936,512]{1,0:T(8,128)} add(%param_0.1523, %mul.2034.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.211 = f32[151936,512]{1,0:T(8,128)} multiply(%add.974.clone.1, %add.974.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1467 = f32[]{:T(128)} constant(0) + %reduce.196 = f32[]{:T(128)} reduce(%square.211, %constant.1467), dimensions={0,1}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.198.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1467), dimensions={0,1}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.150 = (f32[]{:T(128)}, f32[151936,512]{1,0:T(8,128)}, f32[151936,512]{1,0:T(8,128)}, f32[151936,512]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.196, %add.974.clone.1, %add.978.clone.1, %add.979.clone.1, %reduce.198.clone.1) +} + +%region_40.45 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { %reduce_sum.317 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.318 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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.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.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.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.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.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.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.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) -} - -%region_57.62 (reduce_sum.389: f32[], reduce_sum.393: f32[]) -> f32[] { - %reduce_sum.393 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.389 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - 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} - %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_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} - %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} +%fused_computation.321 (param_0.1539: bf16[151936,512], param_1.1718: bf16[151936,2048], param_2.1323: s32[]) -> f32[] { + %param_1.1718 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(1) + %constant.1236 = s32[]{:T(128)} constant(0) + %param_2.1323 = s32[]{:T(128)S(6)} parameter(2) + %dynamic-slice.135 = bf16[151936,512]{1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1718, %constant.1236, %param_2.1323), dynamic_slice_sizes={151936,512}, metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294965759","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %convert_element_type.1229 = f32[151936,512]{1,0:T(8,128)} convert(%dynamic-slice.135), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_0.1539 = bf16[151936,512]{1,0:T(8,128)(2,1)} parameter(0) + %convert_element_type.1228 = f32[151936,512]{1,0:T(8,128)} convert(%param_0.1539), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %add_any.184 = f32[151936,512]{1,0:T(8,128)} add(%convert_element_type.1229, %convert_element_type.1228), metadata={op_name="jit(train_step)/transpose(jvp())/add_any" stack_frame_id=0} + %square.212 = f32[151936,512]{1,0:T(8,128)} multiply(%add_any.184, %add_any.184), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1483 = f32[]{:T(128)} constant(0) + ROOT %reduce.197 = f32[]{:T(128)} reduce(%square.212, %constant.1483), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } %region_9.12 (reduce_sum.186: f32[], reduce_sum.190: f32[]) -> f32[] { - %reduce_sum.190 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} %reduce_sum.186 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.190 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} 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} - %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} +%fused_computation.328 (param_0.1549: bf16[128,151936], param_1.1725: bf16[128]) -> f32[128] { + %param_0.1549 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1250 = f32[128,151936]{1,0:T(8,128)} convert(%param_0.1549), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.520 = f32[1,128,151936]{2,1,0:T(8,128)} bitcast(%convert_element_type.1250), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1725 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %sub.100 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1725), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.71 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%bitcast.520, %sub.100), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.51 = f32[1,128,151936]{2,1,0:T(8,128)} exponential(%sub.71), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %constant.1494 = f32[]{:T(128)} constant(0) + ROOT %reduce.199 = f32[128]{0:T(128)S(1)} reduce(%exp.51, %constant.1494), dimensions={0,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[] { - %reduce_sum.270 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_54.59 (reduce_sum.389: f32[], reduce_sum.393: f32[]) -> f32[] { + %reduce_sum.389 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.393 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + 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.334 (param_0.1548: bf16[128,151936], param_1.1723: f32[128], param_2.1326: bf16[128], param_3.886: s32[128]) -> f32[128] { + %param_3.886 = s32[128]{0:T(128)S(1)} parameter(3) + %eq.26 = s32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.886), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.51 = s32[1,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.21 = pred[1,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.26, %iota.51), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %param_0.1548 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1263 = f32[128,151936]{1,0:T(8,128)} convert(%param_0.1548), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.532 = f32[1,128,151936]{2,1,0:T(8,128)} bitcast(%convert_element_type.1263), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_2.1326 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %sub.99 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1326), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.92 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%bitcast.532, %sub.99), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1723 = f32[128]{0:T(128)S(1)} parameter(1) + %sub.97 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1723), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.91 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%sub.92, %sub.97), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %constant.1492 = f32[]{:T(128)} constant(0) + %broadcast.906 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1492), dimensions={}, metadata={op_name="broadcast.130"} + %mul.1823 = f32[1,128,151936]{2,1,0:T(8,128)} select(%eq.21, %sub.91, %broadcast.906), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.200 = f32[128]{0:T(128)S(1)} reduce(%mul.1823, %constant.1492), dimensions={0,2}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%region_30.35 (reduce_sum.269: f32[], reduce_sum.270: f32[]) -> f32[] { %reduce_sum.269 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.270 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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.340 (param_0.1543: f32[4,6144,512]) -> f32[] { + %param_0.1543 = f32[4,6144,512]{2,0,1:T(4,128)} parameter(0) + %bitcast.536 = f32[6144,4,512]{2,1,0:T(4,128)} bitcast(%param_0.1543), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.215 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%bitcast.536, %bitcast.536), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1487 = f32[]{:T(128)} constant(0) + ROOT %reduce.201 = f32[]{:T(128)} reduce(%square.215, %constant.1487), dimensions={0,1,2}, to_apply=%region_30.35, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%region_32.37 (reduce_sum.263: f32[], reduce_sum.267: f32[]) -> f32[] { - %reduce_sum.267 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_29.34 (reduce_sum.263: f32[], reduce_sum.267: f32[]) -> f32[] { %reduce_sum.263 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.267 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.268 = f32[]{:T(128)} add(%reduce_sum.263, %reduce_sum.267), 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":[]}} } -%region_31.36 (reduce_sum.260: f32[], reduce_sum.261: f32[]) -> f32[] { - %reduce_sum.261 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_28.33 (reduce_sum.260: f32[], reduce_sum.261: f32[]) -> f32[] { %reduce_sum.260 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.261 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - ROOT %tuple.167 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.182, %reduce.183.clone.1) +%fused_computation.342 (param_0.1544: f32[4,512,6144], param_1.1720: f32[4,512,6144]) -> (f32[], f32[]) { + %param_0.1544 = f32[4,512,6144]{2,0,1:T(4,128)} parameter(0) + %bitcast.540 = f32[512,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1544), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.218 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.540, %bitcast.540), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1488 = f32[]{:T(128)} constant(0) + %reduce.202 = f32[]{:T(128)} reduce(%square.218, %constant.1488), dimensions={0,1,2}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1720 = f32[4,512,6144]{2,0,1:T(4,128)} parameter(1) + %bitcast.544.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1720), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.221.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.544.clone.1, %bitcast.544.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.203.clone.1 = f32[]{:T(128)} reduce(%square.221.clone.1, %constant.1488), dimensions={0,1,2}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.173 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.202, %reduce.203.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.345 (param_0.1013: f32[6144,4,512]) -> bf16[4,6144,512] { + %param_0.1013 = f32[6144,4,512]{2,1,0:T(4,128)S(1)} parameter(0) + %copy.209 = bf16[6144,4,512]{2,0,1:T(8,128)(2,1)} copy(%param_0.1013), sharding={devices=[1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\'].value"} + ROOT %bitcast.545 = bf16[4,6144,512]{2,1,0:T(8,128)(2,1)} bitcast(%copy.209), metadata={op_name="jit(train_step)/jvp()/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.346 (param_0.1015: f32[512,4,6144]) -> bf16[4,512,6144] { + %param_0.1015 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.210 = bf16[512,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.1015), sharding={devices=[4,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\'].value"} + ROOT %bitcast.546 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.210), metadata={op_name="jit(train_step)/jvp()/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.347 (param_0.1017: f32[512,4,6144]) -> bf16[4,512,6144] { + %param_0.1017 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.211 = bf16[512,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.1017), sharding={devices=[4,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\'].value"} + ROOT %bitcast.547 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.211), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} } -%region_62.67 (reduce_sum.416: f32[], reduce_sum.417: f32[]) -> f32[] { - %reduce_sum.417 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_59.64 (reduce_sum.416: f32[], reduce_sum.417: f32[]) -> f32[] { %reduce_sum.416 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.417 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.421 = f32[]{:T(128)} add(%reduce_sum.416, %reduce_sum.417), 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":[]}} } -%region_47.52 (reduce_sum.338: f32[], reduce_sum.339: f32[]) -> f32[] { - %reduce_sum.339 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_44.49 (reduce_sum.338: f32[], reduce_sum.339: f32[]) -> f32[] { %reduce_sum.338 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.339 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %add.987.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1147.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - 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) -} - -%region_61.66 (reduce_sum.410: f32[], reduce_sum.414: f32[]) -> f32[] { - %reduce_sum.414 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.348 (param_0.1532: f32[6144,4,512], param_1.1713: f32[], param_2.1319: f32[], param_3.881: f32[], param_4.519: f32[6144,4,512], param_5.463: f32[], param_6.325: f32[4,6144,512], param_7.199: pred[], param_8.121: f32[6144,4,512]) -> (f32[], f32[6144,4,512], f32[6144,4,512], f32[6144,4,512], f32[]) { + %param_0.1532 = f32[6144,4,512]{2,1,0:T(4,128)} parameter(0) + %param_3.881 = f32[]{:T(128)S(6)} parameter(3) + %mul.2080.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%param_3.881), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.199 = pred[]{:T(512)S(6)} parameter(7) + %select_n.308.clone.1 = pred[6144,4,512]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.199), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.325 = f32[4,6144,512]{2,0,1:T(4,128)S(1)} parameter(6) + %bitcast.701.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} bitcast(%param_6.325), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.463 = f32[]{:T(128)} parameter(5) + %div.960.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%param_5.463), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.959.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} divide(%bitcast.701.clone.1, %div.960.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.307.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} select(%select_n.308.clone.1, %bitcast.701.clone.1, %div.959.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1438.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1031.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%constant.1438.clone.1), dimensions={}, metadata={op_name="broadcast.85"} + %mul.2086.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %broadcast.1031.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.121 = f32[6144,4,512]{2,1,0:T(4,128)} parameter(8) + %constant.1442.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.2087.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%constant.1442.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2085.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%param_8.121, %mul.2087.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1007.clone.1 = f32[6144,4,512]{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_2.1319 = f32[]{:T(128)S(6)} parameter(2) + %div.956.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.75.clone.1 = f32[6144,4,512]{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.1441.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.2084.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%constant.1441.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2082.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%integer_pow.75.clone.1, %mul.2084.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.519 = f32[6144,4,512]{2,1,0:T(4,128)} parameter(4) + %constant.1440.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.2083.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%constant.1440.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2081.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%param_4.519, %mul.2083.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1006.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} add(%mul.2082.clone.1, %mul.2081.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1713 = f32[]{:T(128)S(6)} parameter(1) + %div.955.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%param_1.1713), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.954.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} divide(%add.1006.clone.1, %div.955.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.72.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} sqrt(%div.954.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1439.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.1005.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} broadcast(%constant.1439.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.1004.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} add(%sqrt.72.clone.1, %add.1005.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.458.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%div.956.clone.1, %add.1004.clone.1), metadata={op_name="multiply.51"} + %div.953.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} divide(%add.1007.clone.1, %multiply.458.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2079.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%param_0.1532, %broadcast.1031.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1003.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} add(%div.953.clone.1, %mul.2079.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2078.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%mul.2080.clone.1, %add.1003.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1002.clone.1 = f32[6144,4,512]{2,1,0:T(4,128)} add(%param_0.1532, %mul.2078.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.222 = f32[6144,4,512]{2,1,0:T(4,128)} multiply(%add.1002.clone.1, %add.1002.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1476 = f32[]{:T(128)} constant(0) + %reduce.204 = f32[]{:T(128)} reduce(%square.222, %constant.1476), dimensions={0,1,2}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.207.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1476), dimensions={0,1,2}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.151 = (f32[]{:T(128)}, f32[6144,4,512]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.204, %add.1002.clone.1, %add.1006.clone.1, %add.1007.clone.1, %reduce.207.clone.1) +} + +%region_58.63 (reduce_sum.410: f32[], reduce_sum.414: f32[]) -> f32[] { %reduce_sum.410 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.414 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.415 = f32[]{:T(128)} add(%reduce_sum.410, %reduce_sum.414), 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":[]}} } -%region_46.51 (reduce_sum.332: f32[], reduce_sum.333: f32[]) -> f32[] { - %reduce_sum.333 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_43.48 (reduce_sum.332: f32[], reduce_sum.333: f32[]) -> f32[] { %reduce_sum.332 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.333 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %broadcast.887.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1153.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %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} - 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) -} - -%region_60.65 (reduce_sum.407: f32[], reduce_sum.408: f32[]) -> f32[] { - %reduce_sum.408 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.349 (param_0.1533: f32[512,4,6144], param_1.1714: f32[], param_2.1320: f32[], param_3.882: f32[], param_4.520: f32[512,4,6144], param_5.464: f32[], param_6.326: f32[4,512,6144], param_7.200: pred[], param_8.122: f32[512,4,6144]) -> (f32[], f32[512,4,6144], f32[512,4,6144], f32[512,4,6144], f32[]) { + %param_0.1533 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(0) + %param_3.882 = f32[]{:T(128)S(6)} parameter(3) + %mul.2090.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.882), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.200 = pred[]{:T(512)S(6)} parameter(7) + %select_n.312.clone.1 = pred[512,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.200), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.326 = f32[4,512,6144]{2,0,1:T(4,128)} parameter(6) + %bitcast.703.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.326), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.464 = f32[]{:T(128)} parameter(5) + %div.968.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.464), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.967.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} divide(%bitcast.703.clone.1, %div.968.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.311.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} select(%select_n.312.clone.1, %bitcast.703.clone.1, %div.967.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1444.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1037.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1444.clone.1), dimensions={}, metadata={op_name="broadcast.87"} + %mul.2094.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %broadcast.1037.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.122 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(8) + %constant.1448.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.1036.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1448.clone.1), dimensions={}, metadata={op_name="broadcast.86"} + %mul.2093.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%param_8.122, %broadcast.1036.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1012.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%mul.2094.clone.1, %mul.2093.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) + %div.964.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.76.clone.1 = f32[512,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.1447.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.1035.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1447.clone.1), dimensions={}, metadata={op_name="broadcast.75"} + %mul.2092.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.76.clone.1, %broadcast.1035.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.520 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(4) + %constant.1446.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.1034.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1446.clone.1), dimensions={}, metadata={op_name="broadcast.74"} + %mul.2091.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%param_4.520, %broadcast.1034.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1011.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%mul.2092.clone.1, %mul.2091.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1714 = f32[]{:T(128)S(6)} parameter(1) + %div.963.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1714), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.962.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} divide(%add.1011.clone.1, %div.963.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.73.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} sqrt(%div.962.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1445.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.1032.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1445.clone.1), dimensions={}, metadata={op_name="broadcast.67"} + %add.1010.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%sqrt.73.clone.1, %broadcast.1032.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.459.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%div.964.clone.1, %add.1010.clone.1), metadata={op_name="multiply.50"} + %div.961.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} divide(%add.1012.clone.1, %multiply.459.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2089.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1533, %broadcast.1037.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1009.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%div.961.clone.1, %mul.2089.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2088.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%mul.2090.clone.1, %add.1009.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1008.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%param_0.1533, %mul.2088.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.223 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%add.1008.clone.1, %add.1008.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1477 = f32[]{:T(128)} constant(0) + %reduce.205 = f32[]{:T(128)} reduce(%square.223, %constant.1477), dimensions={0,1,2}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.208.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1477), dimensions={0,1,2}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.152 = (f32[]{:T(128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.205, %add.1008.clone.1, %add.1011.clone.1, %add.1012.clone.1, %reduce.208.clone.1) +} + +%region_57.62 (reduce_sum.407: f32[], reduce_sum.408: f32[]) -> f32[] { %reduce_sum.407 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.408 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.409 = f32[]{:T(128)} add(%reduce_sum.407, %reduce_sum.408), 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":[]}} } -%region_45.50 (reduce_sum.326: f32[], reduce_sum.330: f32[]) -> f32[] { - %reduce_sum.330 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_42.47 (reduce_sum.326: f32[], reduce_sum.330: f32[]) -> f32[] { %reduce_sum.326 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.330 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %broadcast.893.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1159.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %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} - 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) -} - -%region_39.44 (reduce_sum.302: f32[], reduce_sum.303: f32[]) -> f32[] { - %reduce_sum.303 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.350 (param_0.1534: f32[512,4,6144], param_1.1715: f32[], param_2.1321: f32[], param_3.883: f32[], param_4.521: f32[512,4,6144], param_5.465: f32[], param_6.327: f32[4,512,6144], param_7.201: pred[], param_8.123: f32[512,4,6144]) -> (f32[], f32[512,4,6144], f32[512,4,6144], f32[512,4,6144], f32[]) { + %param_0.1534 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(0) + %param_3.883 = f32[]{:T(128)S(6)} parameter(3) + %mul.2097.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.883), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.201 = pred[]{:T(512)S(6)} parameter(7) + %select_n.316.clone.1 = pred[512,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.201), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.327 = f32[4,512,6144]{2,0,1:T(4,128)S(1)} parameter(6) + %bitcast.705.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.327), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.465 = f32[]{:T(128)} parameter(5) + %div.976.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.465), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.975.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} divide(%bitcast.705.clone.1, %div.976.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.315.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} select(%select_n.316.clone.1, %bitcast.705.clone.1, %div.975.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1450.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1043.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1450.clone.1), dimensions={}, metadata={op_name="broadcast.87"} + %mul.2101.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%select_n.315.clone.1, %broadcast.1043.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.123 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(8) + %constant.1454.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.1042.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1454.clone.1), dimensions={}, metadata={op_name="broadcast.86"} + %mul.2100.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%param_8.123, %broadcast.1042.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1017.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%mul.2101.clone.1, %mul.2100.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) + %div.972.clone.1 = f32[512,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.77.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%select_n.315.clone.1, %select_n.315.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.1453.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.1041.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1453.clone.1), dimensions={}, metadata={op_name="broadcast.75"} + %mul.2099.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.77.clone.1, %broadcast.1041.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.521 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(4) + %constant.1452.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.1040.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1452.clone.1), dimensions={}, metadata={op_name="broadcast.74"} + %mul.2098.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%param_4.521, %broadcast.1040.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1016.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%mul.2099.clone.1, %mul.2098.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1715 = f32[]{:T(128)S(6)} parameter(1) + %div.971.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1715), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.970.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} divide(%add.1016.clone.1, %div.971.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.74.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} sqrt(%div.970.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1451.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.1038.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1451.clone.1), dimensions={}, metadata={op_name="broadcast.67"} + %add.1015.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%sqrt.74.clone.1, %broadcast.1038.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.460.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%div.972.clone.1, %add.1015.clone.1), metadata={op_name="multiply.49"} + %div.969.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} divide(%add.1017.clone.1, %multiply.460.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2096.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1534, %broadcast.1043.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1014.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%div.969.clone.1, %mul.2096.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2095.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%mul.2097.clone.1, %add.1014.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1013.clone.1 = f32[512,4,6144]{2,1,0:T(4,128)} add(%param_0.1534, %mul.2095.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.224 = f32[512,4,6144]{2,1,0:T(4,128)} multiply(%add.1013.clone.1, %add.1013.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1478 = f32[]{:T(128)} constant(0) + %reduce.206 = f32[]{:T(128)} reduce(%square.224, %constant.1478), dimensions={0,1,2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1478), dimensions={0,1,2}, to_apply=%region_42.47, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.153 = (f32[]{:T(128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.206, %add.1013.clone.1, %add.1016.clone.1, %add.1017.clone.1, %reduce.209.clone.1) +} + +%region_36.41 (reduce_sum.302: f32[], reduce_sum.303: f32[]) -> f32[] { %reduce_sum.302 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.303 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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.367 (param_0.1540: f32[4,512,16,128]) -> f32[] { + %param_0.1540 = f32[4,512,16,128]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.554 = f32[512,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1540), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.227 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%bitcast.554, %bitcast.554), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1484 = f32[]{:T(128)} constant(0) + ROOT %reduce.210 = f32[]{:T(128)} reduce(%square.227, %constant.1484), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%region_38.43 (reduce_sum.296: f32[], reduce_sum.297: f32[]) -> f32[] { - %reduce_sum.297 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_35.40 (reduce_sum.296: f32[], reduce_sum.297: f32[]) -> f32[] { %reduce_sum.296 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.297 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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.369 (param_0.1541: f32[4,16,128,512]) -> f32[] { + %param_0.1541 = f32[4,16,128,512]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.558 = f32[16,4,128,512]{3,2,1,0:T(8,128)} bitcast(%param_0.1541), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.230 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%bitcast.558, %bitcast.558), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1485 = f32[]{:T(128)} constant(0) + ROOT %reduce.211 = f32[]{:T(128)} reduce(%square.230, %constant.1485), dimensions={0,1,2,3}, to_apply=%region_35.40, 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.370 (param_0.1066: f32[16,4,128,512]) -> bf16[4,16,128,512] { + %param_0.1066 = f32[16,4,128,512]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.212 = bf16[16,4,128,512]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.1066), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\'].value"} + ROOT %bitcast.559 = bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.212), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} } -%region_68.73 (reduce_sum.449: f32[], reduce_sum.450: f32[]) -> f32[] { - %reduce_sum.450 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_65.70 (reduce_sum.449: f32[], reduce_sum.450: f32[]) -> f32[] { %reduce_sum.449 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.450 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.451 = f32[]{:T(128)} add(%reduce_sum.449, %reduce_sum.450), 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":[]}} } -%region_53.58 (reduce_sum.368: f32[], reduce_sum.372: f32[]) -> f32[] { - %reduce_sum.372 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_50.55 (reduce_sum.368: f32[], reduce_sum.372: f32[]) -> f32[] { %reduce_sum.368 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.372 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %add.955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1111.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - 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) -} - -%region_67.72 (reduce_sum.443: f32[], reduce_sum.444: f32[]) -> f32[] { - %reduce_sum.444 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.371 (param_0.1526: f32[512,4,16,128], param_1.1707: f32[], param_2.1313: f32[], param_3.875: f32[], param_4.513: f32[512,4,16,128], param_5.457: f32[], param_6.319: f32[4,512,16,128], param_7.193: pred[], param_8.115: f32[512,4,16,128]) -> (f32[], f32[512,4,16,128], f32[512,4,16,128], f32[512,4,16,128], f32[]) { + %param_0.1526 = f32[512,4,16,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %param_3.875 = f32[]{:T(128)S(6)} parameter(3) + %mul.2053.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_3.875), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.193 = pred[]{:T(512)S(6)} parameter(7) + %select_n.296.clone.1 = pred[512,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.193), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.319 = f32[4,512,16,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.695.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.319), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.457 = f32[]{:T(128)} parameter(5) + %div.936.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.457), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.935.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} divide(%bitcast.695.clone.1, %div.936.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.295.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} select(%select_n.296.clone.1, %bitcast.695.clone.1, %div.935.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1420.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1021.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1420.clone.1), dimensions={}, metadata={op_name="broadcast.77"} + %mul.2059.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.295.clone.1, %broadcast.1021.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.115 = f32[512,4,16,128]{3,2,1,0:T(8,128)} parameter(8) + %constant.1424.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.2060.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1424.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2058.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.115, %mul.2060.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.990.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} add(%mul.2059.clone.1, %mul.2058.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1313 = f32[]{:T(128)S(6)} parameter(2) + %div.932.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1313), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.72.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,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.1423.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.2057.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1423.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2055.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.72.clone.1, %mul.2057.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.513 = f32[512,4,16,128]{3,2,1,0:T(8,128)} parameter(4) + %constant.1422.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.2056.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1422.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2054.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.513, %mul.2056.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.989.clone.1 = f32[512,4,16,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.1707 = f32[]{:T(128)S(6)} parameter(1) + %div.931.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1707), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.930.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} divide(%add.989.clone.1, %div.931.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.69.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} sqrt(%div.930.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1421.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.988.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1421.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.987.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} add(%sqrt.69.clone.1, %add.988.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.455.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%div.932.clone.1, %add.987.clone.1), metadata={op_name="multiply.57"} + %div.929.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} divide(%add.990.clone.1, %multiply.455.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2052.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1526, %broadcast.1021.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.986.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} add(%div.929.clone.1, %mul.2052.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2051.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%mul.2053.clone.1, %add.986.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.985.clone.1 = f32[512,4,16,128]{3,2,1,0:T(8,128)S(1)} add(%param_0.1526, %mul.2051.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.231 = f32[512,4,16,128]{3,2,1,0:T(8,128)} multiply(%add.985.clone.1, %add.985.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1470 = f32[]{:T(128)} constant(0) + %reduce.212 = f32[]{:T(128)} reduce(%square.231, %constant.1470), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.214.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1470), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.154 = (f32[]{:T(128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)S(1)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.212, %add.985.clone.1, %add.989.clone.1, %add.990.clone.1, %reduce.214.clone.1) +} + +%region_64.69 (reduce_sum.443: f32[], reduce_sum.444: f32[]) -> f32[] { %reduce_sum.443 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.444 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.445 = f32[]{:T(128)} add(%reduce_sum.443, %reduce_sum.444), 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":[]}} } -%region_52.57 (reduce_sum.365: f32[], reduce_sum.366: f32[]) -> f32[] { - %reduce_sum.366 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_49.54 (reduce_sum.365: f32[], reduce_sum.366: f32[]) -> f32[] { %reduce_sum.365 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.366 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %add.961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1117.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - 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) -} - -%region_41.46 (reduce_sum.311: f32[], reduce_sum.312: f32[]) -> f32[] { - %reduce_sum.312 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.372 (param_0.1527: f32[16,4,128,512], param_1.1708: f32[], param_2.1314: f32[], param_3.876: f32[], param_4.514: f32[16,4,128,512], param_5.458: f32[], param_6.320: f32[4,16,128,512], param_7.194: pred[], param_8.116: f32[16,4,128,512]) -> (f32[], f32[16,4,128,512], f32[16,4,128,512], f32[16,4,128,512], f32[]) { + %param_0.1527 = f32[16,4,128,512]{3,2,1,0:T(8,128)} parameter(0) + %param_3.876 = f32[]{:T(128)S(6)} parameter(3) + %mul.2063.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} broadcast(%param_3.876), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.194 = pred[]{:T(512)S(6)} parameter(7) + %select_n.300.clone.1 = pred[16,4,128,512]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.194), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.320 = f32[4,16,128,512]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.697.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} bitcast(%param_6.320), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.458 = f32[]{:T(128)} parameter(5) + %div.944.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} broadcast(%param_5.458), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.943.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} divide(%bitcast.697.clone.1, %div.944.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.299.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} select(%select_n.300.clone.1, %bitcast.697.clone.1, %div.943.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1426.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1023.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} broadcast(%constant.1426.clone.1), dimensions={}, metadata={op_name="broadcast.78"} + %mul.2069.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%select_n.299.clone.1, %broadcast.1023.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.116 = f32[16,4,128,512]{3,2,1,0:T(8,128)} parameter(8) + %constant.1430.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.2070.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} broadcast(%constant.1430.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2068.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%param_8.116, %mul.2070.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.996.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} add(%mul.2069.clone.1, %mul.2068.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) + %div.940.clone.1 = f32[16,4,128,512]{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.73.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,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.1429.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.2067.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} broadcast(%constant.1429.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2065.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%integer_pow.73.clone.1, %mul.2067.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.514 = f32[16,4,128,512]{3,2,1,0:T(8,128)} parameter(4) + %constant.1428.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.2066.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} broadcast(%constant.1428.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2064.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%param_4.514, %mul.2066.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.995.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} add(%mul.2065.clone.1, %mul.2064.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1708 = f32[]{:T(128)S(6)} parameter(1) + %div.939.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} broadcast(%param_1.1708), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.938.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} divide(%add.995.clone.1, %div.939.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.70.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} sqrt(%div.938.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1427.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.994.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} broadcast(%constant.1427.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.993.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} add(%sqrt.70.clone.1, %add.994.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.456.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%div.940.clone.1, %add.993.clone.1), metadata={op_name="multiply.56"} + %div.937.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} divide(%add.996.clone.1, %multiply.456.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2062.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%param_0.1527, %broadcast.1023.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.992.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} add(%div.937.clone.1, %mul.2062.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2061.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%mul.2063.clone.1, %add.992.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.991.clone.1 = f32[16,4,128,512]{3,2,1,0:T(8,128)} add(%param_0.1527, %mul.2061.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.232 = f32[16,4,128,512]{3,2,1,0:T(8,128)} multiply(%add.991.clone.1, %add.991.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1471 = f32[]{:T(128)} constant(0) + %reduce.213 = f32[]{:T(128)} reduce(%square.232, %constant.1471), dimensions={0,1,2,3}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.215.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1471), dimensions={0,1,2,3}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.155 = (f32[]{:T(128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.213, %add.991.clone.1, %add.995.clone.1, %add.996.clone.1, %reduce.215.clone.1) +} + +%region_38.43 (reduce_sum.311: f32[], reduce_sum.312: f32[]) -> f32[] { %reduce_sum.311 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.312 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.316 = f32[]{:T(128)} add(%reduce_sum.311, %reduce_sum.312), 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":[]}} } -%region_36.41 (reduce_sum.284: f32[], reduce_sum.288: f32[]) -> f32[] { - %reduce_sum.288 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_33.38 (reduce_sum.284: f32[], reduce_sum.288: f32[]) -> f32[] { %reduce_sum.284 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.288 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - ROOT %tuple.168 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.196, %reduce.197.clone.1) +%fused_computation.384 (param_0.1542: f32[4,512,8,128], param_1.1719: f32[4,512,8,128]) -> (f32[], f32[]) { + %param_0.1542 = f32[4,512,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.563 = f32[512,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1542), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.235 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.563, %bitcast.563), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1486 = f32[]{:T(128)} constant(0) + %reduce.216 = f32[]{:T(128)} reduce(%square.235, %constant.1486), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1719 = f32[4,512,8,128]{3,2,0,1:T(8,128)S(1)} parameter(1) + %bitcast.567.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1719), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.238.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.567.clone.1, %bitcast.567.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.217.clone.1 = f32[]{:T(128)} reduce(%square.238.clone.1, %constant.1486), dimensions={0,1,2,3}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.174 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.216, %reduce.217.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.387 (param_0.1098: f32[512,4,8,128]) -> bf16[4,512,8,128] { + %param_0.1098 = f32[512,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.213 = bf16[512,4,8,128]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1098), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\'].value"} + ROOT %bitcast.568 = bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.213), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} } -%region_70.75 (reduce_sum.458: f32[], reduce_sum.459: f32[]) -> f32[] { - %reduce_sum.459 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_67.72 (reduce_sum.458: f32[], reduce_sum.459: f32[]) -> f32[] { %reduce_sum.458 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.459 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.463 = f32[]{:T(128)} add(%reduce_sum.458, %reduce_sum.459), 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":[]}} } -%region_55.60 (reduce_sum.380: f32[], reduce_sum.381: f32[]) -> f32[] { - %reduce_sum.381 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_52.57 (reduce_sum.380: f32[], reduce_sum.381: f32[]) -> f32[] { %reduce_sum.380 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.381 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %broadcast.845.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1099.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %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} - 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) -} - -%region_65.70 (reduce_sum.431: f32[], reduce_sum.435: f32[]) -> f32[] { - %reduce_sum.435 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.388 (param_0.1524: f32[512,4,8,128], param_1.1705: f32[], param_2.1311: f32[], param_3.873: f32[], param_4.511: f32[512,4,8,128], param_5.455: f32[], param_6.317: f32[4,512,8,128], param_7.191: pred[], param_8.113: f32[512,4,8,128]) -> (f32[], f32[512,4,8,128], f32[512,4,8,128], f32[512,4,8,128], f32[]) { + %param_0.1524 = f32[512,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.873 = f32[]{:T(128)S(6)} parameter(3) + %mul.2046.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.873), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.191 = pred[]{:T(512)S(6)} parameter(7) + %select_n.292.clone.1 = pred[512,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.317 = f32[4,512,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.693.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.317), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.455 = f32[]{:T(128)} parameter(5) + %div.928.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.455), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.927.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.693.clone.1, %div.928.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.291.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.292.clone.1, %bitcast.693.clone.1, %div.927.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1414.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1019.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1414.clone.1), dimensions={}, metadata={op_name="broadcast.82"} + %mul.2050.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %broadcast.1019.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.113 = f32[512,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %constant.1418.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.1018.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1418.clone.1), dimensions={}, metadata={op_name="broadcast.81"} + %mul.2049.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.113, %broadcast.1018.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.984.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2050.clone.1, %mul.2049.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1311 = f32[]{:T(128)S(6)} parameter(2) + %div.924.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1311), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.71.clone.1 = f32[512,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.1417.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.1017.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1417.clone.1), dimensions={}, metadata={op_name="broadcast.71"} + %mul.2048.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.71.clone.1, %broadcast.1017.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.511 = f32[512,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %constant.1416.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.1016.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1416.clone.1), dimensions={}, metadata={op_name="broadcast.70"} + %mul.2047.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.511, %broadcast.1016.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.983.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2048.clone.1, %mul.2047.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1705 = f32[]{:T(128)S(6)} parameter(1) + %div.923.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1705), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.922.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} divide(%add.983.clone.1, %div.923.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.68.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.922.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1415.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.1014.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1415.clone.1), dimensions={}, metadata={op_name="broadcast.65"} + %add.982.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.68.clone.1, %broadcast.1014.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.454.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.924.clone.1, %add.982.clone.1), metadata={op_name="multiply.59"} + %div.921.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} divide(%add.984.clone.1, %multiply.454.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2045.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1524, %broadcast.1019.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.981.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%div.921.clone.1, %mul.2045.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2044.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.2046.clone.1, %add.981.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.980.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1524, %mul.2044.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.239 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.980.clone.1, %add.980.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1468 = f32[]{:T(128)} constant(0) + %reduce.218 = f32[]{:T(128)} reduce(%square.239, %constant.1468), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.220.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1468), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.156 = (f32[]{:T(128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.218, %add.980.clone.1, %add.983.clone.1, %add.984.clone.1, %reduce.220.clone.1) +} + +%region_62.67 (reduce_sum.431: f32[], reduce_sum.435: f32[]) -> f32[] { %reduce_sum.431 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.435 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.436 = f32[]{:T(128)} add(%reduce_sum.431, %reduce_sum.435), 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":[]}} } -%region_50.55 (reduce_sum.353: f32[], reduce_sum.354: f32[]) -> f32[] { - %reduce_sum.354 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_47.52 (reduce_sum.353: f32[], reduce_sum.354: f32[]) -> f32[] { %reduce_sum.353 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.354 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %broadcast.867.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1129.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %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) -} - -%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.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} +%fused_computation.389 (param_0.1529: f32[512,4,8,128], param_1.1710: f32[], param_2.1316: f32[], param_3.878: f32[], param_4.516: f32[512,4,8,128], param_5.460: f32[], param_6.322: f32[4,512,8,128], param_7.196: pred[], param_8.118: f32[512,4,8,128]) -> (f32[], f32[512,4,8,128], f32[512,4,8,128], f32[512,4,8,128], f32[]) { + %param_0.1529 = f32[512,4,8,128]{3,2,1,0:T(8,128)} parameter(0) + %param_3.878 = f32[]{:T(128)S(6)} parameter(3) + %mul.2073.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.878), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.196 = pred[]{:T(512)S(6)} parameter(7) + %select_n.304.clone.1 = pred[512,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.196), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.322 = f32[4,512,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) + %bitcast.699.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.322), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.460 = f32[]{:T(128)} parameter(5) + %div.952.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.460), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.951.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.699.clone.1, %div.952.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.303.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.304.clone.1, %bitcast.699.clone.1, %div.951.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1432.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1029.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1432.clone.1), dimensions={}, metadata={op_name="broadcast.82"} + %mul.2077.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.303.clone.1, %broadcast.1029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.118 = f32[512,4,8,128]{3,2,1,0:T(8,128)} parameter(8) + %constant.1436.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.1028.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1436.clone.1), dimensions={}, metadata={op_name="broadcast.81"} + %mul.2076.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.118, %broadcast.1028.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1001.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2077.clone.1, %mul.2076.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) + %div.948.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.74.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,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.1435.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.1027.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1435.clone.1), dimensions={}, metadata={op_name="broadcast.71"} + %mul.2075.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.74.clone.1, %broadcast.1027.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.516 = f32[512,4,8,128]{3,2,1,0:T(8,128)} parameter(4) + %constant.1434.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.1026.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1434.clone.1), dimensions={}, metadata={op_name="broadcast.70"} + %mul.2074.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.516, %broadcast.1026.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.1000.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%mul.2075.clone.1, %mul.2074.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1710 = f32[]{:T(128)S(6)} parameter(1) + %div.947.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1710), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.946.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} divide(%add.1000.clone.1, %div.947.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.71.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.946.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1433.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.1024.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1433.clone.1), dimensions={}, metadata={op_name="broadcast.65"} + %add.999.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.71.clone.1, %broadcast.1024.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.457.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.948.clone.1, %add.999.clone.1), metadata={op_name="multiply.54"} + %div.945.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} divide(%add.1001.clone.1, %multiply.457.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2072.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1529, %broadcast.1029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.998.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%div.945.clone.1, %mul.2072.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2071.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.2073.clone.1, %add.998.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.997.clone.1 = f32[512,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1529, %mul.2071.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.240 = f32[512,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.997.clone.1, %add.997.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1473 = f32[]{:T(128)} constant(0) + %reduce.219 = f32[]{:T(128)} reduce(%square.240, %constant.1473), dimensions={0,1,2,3}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1473), dimensions={0,1,2,3}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.157 = (f32[]{:T(128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.219, %add.997.clone.1, %add.1000.clone.1, %add.1001.clone.1, %reduce.221.clone.1) +} + +%fused_computation.405 (param_0.1174: bf16[1,128,2048], param_1.1238: f32[128], param_2.804: f32[128], param_3.435: bf16[1,128,2048], param_4.244: bf16[2048]) -> bf16[1,128,2048] { + %param_3.435 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %dot_general.478 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.244), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %dot_general.464 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.435, %dot_general.478), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0} + %convert_element_type.1281 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.464), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_2.804 = f32[128]{0:T(128)S(1)} parameter(2) + %mul.1903 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.804), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1895 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1281, %mul.1903), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1174 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1292 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1174), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1238 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1902 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1238), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1901 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1292, %mul.1902), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %add_any.193 = f32[1,128,2048]{2,1,0:T(8,128)} add(%mul.1895, %mul.1901), metadata={op_name="jit(train_step)/transpose(jvp())/add_any" stack_frame_id=0} + ROOT %convert_element_type.1279 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.193), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} } %region_7.10 (reduce_sum.171: f32[], reduce_sum.184: f32[]) -> f32[] { - %reduce_sum.184 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - %reduce_sum.171 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - 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":[]}} + %reduce_sum.171 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.184 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.185 = f32[]{:T(128)} add(%reduce_sum.171, %reduce_sum.184), 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.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.406 (param_0.1550: bf16[1,128,2048]) -> f32[128] { + %param_0.1550 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1283 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1550), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %square.243 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1283, %convert_element_type.1283), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.1495 = f32[]{:T(128)} constant(0) + ROOT %reduce.222 = f32[128]{0:T(128)S(1)} reduce(%square.243, %constant.1495), dimensions={0,2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_12.15 (reduce_sum.198: f32[], reduce_sum.199: f32[]) -> f32[] { - %reduce_sum.199 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - %reduce_sum.198 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - 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} -} - -%region_10.13 (dot_general.190: bf16[], dot_general.191: bf16[]) -> bf16[] { - %dot_general.191 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general"} - %dot_general.190 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general"} - 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.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.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.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.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.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.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.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} - %cos.41.clone.1 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.717), metadata={op_name="jit(train_step)/layers/cos" stack_frame_id=0} - %convert_element_type.1391.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.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) - %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} - 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) - %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} - 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} -} - -%region_35.40 (reduce_sum.281: f32[], reduce_sum.282: f32[]) -> f32[] { - %reduce_sum.282 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.198 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum"} + %reduce_sum.199 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum"} + ROOT %reduce_sum.200 = f32[]{:T(128)} add(%reduce_sum.198, %reduce_sum.199), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.408 (param_0.1546: bf16[1,128,2048], param_1.1721: bf16[1,128,2048], param_2.1324: bf16[2048]) -> f32[128] { + %param_0.1546 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1290 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1546), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1721 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1324 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.477 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1324), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %dot_general.463 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1721, %dot_general.477), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0} + %convert_element_type.1289 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.463), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %mul.1899 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1290, %convert_element_type.1289), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.1490 = f32[]{:T(128)} constant(0) + ROOT %reduce.223 = f32[128]{0:T(128)S(1)} reduce(%mul.1899, %constant.1490), dimensions={0,2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0} +} + +%fused_computation.410 (param_0.1189: bf16[2048], param_1.1256: f32[128], param_2.812: bf16[1,128,2048]) -> bf16[128,2048] { + %param_0.1189 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.476 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1189), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %param_2.812 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1298 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_2.812), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1256 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1910 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1256), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1909 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1298, %mul.1910), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.1297 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1909), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %dot_general.468 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.476, %convert_element_type.1297), metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + ROOT %bitcast.595 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%dot_general.468) +} + +%fused_computation.363.clone (param_0.1399: bf16[128,151936], param_1.1581: f32[128], param_2.1127: f32[128], param_3.702: bf16[128], param_4.332: s32[128], param_5.278: f32[128]) -> bf16[128,151936] { + %param_5.278 = f32[128]{0:T(128)S(1)} parameter(5) + %mul.1989 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.278), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_1.1581 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1988 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1581), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1399 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1329 = f32[128,151936]{1,0:T(8,128)} convert(%param_0.1399), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.675 = f32[1,128,151936]{2,1,0:T(8,128)} bitcast(%convert_element_type.1329), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_3.702 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(3) + %sub.103 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.702), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.102 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%bitcast.675, %sub.103), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.62 = f32[1,128,151936]{2,1,0:T(8,128)} exponential(%sub.102), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %mul.1987 = f32[1,128,151936]{2,1,0:T(8,128)} multiply(%mul.1988, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1127 = f32[128]{0:T(128)S(1)} parameter(2) + %div.870 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1127), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.869 = f32[1,128,151936]{2,1,0:T(8,128)} divide(%mul.1987, %div.870), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_4.332 = s32[128]{0:T(128)S(1)} parameter(4) + %eq.29 = s32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.332), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.57 = s32[1,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.28 = pred[1,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.29, %iota.57), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %convert_element_type.1328 = f32[1,128,151936]{2,1,0:T(8,128)} convert(%eq.28), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.101 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%div.869, %convert_element_type.1328), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1986 = f32[1,128,151936]{2,1,0:T(8,128)} multiply(%mul.1989, %sub.101), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %convert_element_type.1327 = bf16[1,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.1986), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + ROOT %bitcast.674 = bf16[128,151936]{1,0:T(8,128)(2,1)} bitcast(%convert_element_type.1327) +} + +%fused_computation.409 (param_0.1400: bf16[2048], param_1.1582: bf16[1,128,2048], param_2.1128: f32[128], param_3.703: bf16[128,151936], param_4.333: f32[128], param_5.279: f32[128], param_6.154: bf16[128], param_7.87: s32[128], param_8.5: f32[128]) -> bf16[151936,2048] { + %param_3.703 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.333 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.279 = f32[128]{0:T(128)S(1)} parameter(5) + %param_6.154 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %param_7.87 = s32[128]{0:T(128)S(1)} parameter(7) + %param_8.5 = f32[128]{0:T(128)S(1)} parameter(8) + %fusion.340 = bf16[128,151936]{1,0:T(8,128)(2,1)} fusion(%param_3.703, %param_4.333, %param_5.279, %param_6.154, %param_7.87, /*index=5*/%param_8.5), kind=kLoop, calls=%fused_computation.363.clone + %param_0.1400 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %param_2.1128 = f32[128]{0:T(128)S(1)} parameter(2) + %param_1.1582 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.317 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_0.1400, %param_2.1128, %param_1.1582), kind=kLoop, calls=%fused_computation.410 + ROOT %convolution.114 = bf16[151936,2048]{1,0:T(8,128)(2,1)} convolution(%fusion.340, %fusion.317), dim_labels=fb_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} +} + +%region_10.13 (dot_general.204: bf16[], dot_general.205: bf16[]) -> bf16[] { + %dot_general.204 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + %dot_general.205 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + ROOT %add.416 = bf16[]{:T(256)} add(%dot_general.204, %dot_general.205), metadata={op_name="add.79"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.363.clone.1.clone.clone (param_0.1513: bf16[128,151936], param_1.1698: f32[128], param_2.1289: f32[128], param_3.865: bf16[128], param_4.500: s32[128], param_5.434: f32[128]) -> bf16[128,151936] { + %param_5.434 = f32[128]{0:T(128)S(1)} parameter(5) + %mul.2147 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.434), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_1.1698 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2146 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1698), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1513 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1370 = f32[128,151936]{1,0:T(8,128)} convert(%param_0.1513), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.753 = f32[1,128,151936]{2,1,0:T(8,128)} bitcast(%convert_element_type.1370), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_3.865 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(3) + %sub.114 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.865), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.113 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%bitcast.753, %sub.114), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.67 = f32[1,128,151936]{2,1,0:T(8,128)} exponential(%sub.113), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %mul.2145 = f32[1,128,151936]{2,1,0:T(8,128)} multiply(%mul.2146, %exp.67), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1289 = f32[128]{0:T(128)S(1)} parameter(2) + %div.982 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1289), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.981 = f32[1,128,151936]{2,1,0:T(8,128)} divide(%mul.2145, %div.982), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_4.500 = s32[128]{0:T(128)S(1)} parameter(4) + %eq.35 = s32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.500), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.60 = s32[1,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.34 = pred[1,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.35, %iota.60), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %convert_element_type.1369 = f32[1,128,151936]{2,1,0:T(8,128)} convert(%eq.34), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.112 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%div.981, %convert_element_type.1369), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2144 = f32[1,128,151936]{2,1,0:T(8,128)} multiply(%mul.2147, %sub.112), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %convert_element_type.1368 = bf16[1,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2144), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + ROOT %bitcast.752 = bf16[128,151936]{1,0:T(8,128)(2,1)} bitcast(%convert_element_type.1368) +} + +%bitcast_fusion.2 (bitcast_input.2: bf16[151936,2048]) -> bf16[151936,2048] { + %bitcast_input.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.760 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.2) +} + +%fused_computation.413 (param_0.1512: f32[128], param_1.1697: bf16[1,128,2048], param_2.1290: bf16[151936,2048], param_3.866: bf16[128,151936], param_4.501: f32[128], param_5.435: f32[128], param_6.299: bf16[128], param_7.189: s32[128], param_8.111: f32[128]) -> (bf16[2048], bf16[1,128,2048]) { + %param_3.866 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.501 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.435 = f32[128]{0:T(128)S(1)} parameter(5) + %param_6.299 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %param_7.189 = s32[128]{0:T(128)S(1)} parameter(7) + %param_8.111 = f32[128]{0:T(128)S(1)} parameter(8) + %fusion.341.clone.1 = bf16[128,151936]{1,0:T(8,128)(2,1)} fusion(%param_3.866, %param_4.501, %param_5.435, %param_6.299, %param_7.189, /*index=5*/%param_8.111), kind=kLoop, calls=%fused_computation.363.clone.1.clone.clone + %param_2.1290 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) + %fusion.356 = bf16[151936,2048]{1,0:T(8,128)(2,1)} fusion(%param_2.1290), kind=kLoop, calls=%bitcast_fusion.2 + %convolution.113.clone.1 = bf16[128,2048]{1,0:T(8,128)(2,1)} convolution(%fusion.341.clone.1, %fusion.356), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} + %bitcast.593.clone.1 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.113.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} + %param_1.1697 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1302 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_1.1697), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_0.1512 = f32[128]{0:T(128)S(1)} parameter(0) + %mul.1914 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1512), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1913 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1302, %mul.1914), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.1301 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1913), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %multiply.436 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.593.clone.1, %convert_element_type.1301), metadata={op_name="multiply.354"} + %constant.1318 = bf16[]{:T(256)} constant(0) + %reduce.224 = bf16[2048]{0:T(1024)(128)(2,1)} reduce(%multiply.436, %constant.1318), dimensions={0,1}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0} + ROOT %tuple.172 = (bf16[2048]{0:T(1024)(128)(2,1)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.224, %bitcast.593.clone.1) +} + +%fused_computation.421 (param_0.1238: f32[128], param_1.1308: f32[64]) -> (bf16[1,128,1,64], bf16[1,128,1,64]) { + %param_0.1238 = f32[128]{0:T(128)S(1)} parameter(0) + %div.732 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_0.1238), dimensions={1}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1308 = f32[64]{0:T(128)S(1)} parameter(1) + %div.738 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_1.1308), dimensions={3}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.729 = f32[1,128,1,64]{3,1,2,0:T(8,128)} divide(%div.732, %div.738), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sin.38 = f32[1,128,1,64]{3,1,2,0:T(8,128)} sine(%div.729), metadata={op_name="jit(train_step)/sin" stack_frame_id=0} + %convert_element_type.1310 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %cos.41.clone.1 = f32[1,128,1,64]{3,1,2,0:T(8,128)} cosine(%div.729), metadata={op_name="jit(train_step)/cos" stack_frame_id=0} + %convert_element_type.1309.clone.1 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + ROOT %tuple.167 = (bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1310, %convert_element_type.1309.clone.1) +} + +%fused_computation.424 (param_0.1209: bf16[1,128,1,64]) -> bf16[1,128,1,128] { + %param_0.1209 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1310 = bf16[]{:T(256)} constant(-inf) + %pad.54 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1209, %constant.1310), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %pad.53 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1209, %constant.1310), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + ROOT %maximum.46 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} maximum(%pad.54, %pad.53), metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} +} + +%fused_computation.425 (param_0.1211: bf16[1,128,1,64]) -> bf16[1,128,1,128] { + %param_0.1211 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1309 = bf16[]{:T(256)} constant(-inf) + %pad.56 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1211, %constant.1309), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %pad.55 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1211, %constant.1309), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + ROOT %maximum.47 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} maximum(%pad.56, %pad.55), metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} +} + +%region_32.37 (reduce_sum.281: f32[], reduce_sum.282: f32[]) -> f32[] { %reduce_sum.281 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.282 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.283 = f32[]{:T(128)} add(%reduce_sum.281, %reduce_sum.282), 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":[]}} } -%region_34.39 (reduce_sum.275: f32[], reduce_sum.276: f32[]) -> f32[] { - %reduce_sum.276 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.427 (param_0.1537: f32[4,2048]) -> f32[] { + %param_0.1537 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.601 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1537), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.246 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.601, %bitcast.601), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1481 = f32[]{:T(128)} constant(0) + ROOT %reduce.225 = f32[]{:T(128)} reduce(%square.246, %constant.1481), dimensions={0,1}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} +} + +%region_31.36 (reduce_sum.275: f32[], reduce_sum.276: f32[]) -> f32[] { %reduce_sum.275 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.276 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - ROOT %tuple.169 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.205, %reduce.206.clone.1) +%fused_computation.429 (param_0.1538: f32[4,2048]) -> f32[] { + %param_0.1538 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) + %bitcast.605 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1538), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.249 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.605, %bitcast.605), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1482 = f32[]{:T(128)} constant(0) + ROOT %reduce.226 = f32[]{:T(128)} reduce(%square.249, %constant.1482), dimensions={0,1}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%region_64.69 (reduce_sum.428: f32[], reduce_sum.429: f32[]) -> f32[] { - %reduce_sum.429 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_61.66 (reduce_sum.428: f32[], reduce_sum.429: f32[]) -> f32[] { %reduce_sum.428 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.429 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.430 = f32[]{:T(128)} add(%reduce_sum.428, %reduce_sum.429), 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":[]}} } -%region_49.54 (reduce_sum.347: f32[], reduce_sum.351: f32[]) -> f32[] { - %reduce_sum.351 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_46.51 (reduce_sum.347: f32[], reduce_sum.351: f32[]) -> f32[] { %reduce_sum.347 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.351 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %broadcast.873.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1135.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %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} - 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) -} - -%region_63.68 (reduce_sum.422: f32[], reduce_sum.423: f32[]) -> f32[] { - %reduce_sum.423 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.430 (param_0.1530: f32[2048,4], param_1.1711: f32[], param_2.1317: f32[], param_3.879: f32[], param_4.517: f32[2048,4], param_5.461: f32[], param_6.323: f32[4,2048], param_7.197: pred[], param_8.119: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1530 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.879 = f32[]{:T(128)S(6)} parameter(3) + %mul.2012.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.879), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.197 = pred[]{:T(512)S(6)} parameter(7) + %select_n.276.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.323 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.689.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.323), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.461 = f32[]{:T(128)} parameter(5) + %div.896.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.461), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.895.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.689.clone.1, %div.896.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.275.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.276.clone.1, %bitcast.689.clone.1, %div.895.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1388.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1003.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1388.clone.1), dimensions={}, metadata={op_name="broadcast.84"} + %mul.2016.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %broadcast.1003.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.119 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) + %constant.1392.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.1002.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1392.clone.1), dimensions={}, metadata={op_name="broadcast.83"} + %mul.2015.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.119, %broadcast.1002.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.962.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2016.clone.1, %mul.2015.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) + %div.892.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.67.clone.1 = f32[2048,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.1391.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.1001.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1391.clone.1), dimensions={}, metadata={op_name="broadcast.73"} + %mul.2014.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.1001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.517 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %constant.1390.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.1000.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1390.clone.1), dimensions={}, metadata={op_name="broadcast.72"} + %mul.2013.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.517, %broadcast.1000.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.961.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.2014.clone.1, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1711 = f32[]{:T(128)S(6)} parameter(1) + %div.891.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1711), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.890.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.961.clone.1, %div.891.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.64.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.890.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1389.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.998.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1389.clone.1), dimensions={}, metadata={op_name="broadcast.66"} + %add.960.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.998.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.450.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.892.clone.1, %add.960.clone.1), metadata={op_name="multiply.53"} + %div.889.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.962.clone.1, %multiply.450.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2011.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1530, %broadcast.1003.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.959.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.889.clone.1, %mul.2011.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2010.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.2012.clone.1, %add.959.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.958.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1530, %mul.2010.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.250 = f32[2048,4]{0,1:T(4,128)} multiply(%add.958.clone.1, %add.958.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1474 = f32[]{:T(128)} constant(0) + %reduce.227 = f32[]{:T(128)} reduce(%square.250, %constant.1474), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.229.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1474), dimensions={0,1}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.158 = (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.227, %add.958.clone.1, %add.961.clone.1, %add.962.clone.1, %reduce.229.clone.1) +} + +%region_60.65 (reduce_sum.422: f32[], reduce_sum.423: f32[]) -> f32[] { %reduce_sum.422 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.423 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.424 = f32[]{:T(128)} add(%reduce_sum.422, %reduce_sum.423), 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":[]}} } -%region_48.53 (reduce_sum.344: f32[], reduce_sum.345: f32[]) -> f32[] { - %reduce_sum.345 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_45.50 (reduce_sum.344: f32[], reduce_sum.345: f32[]) -> f32[] { %reduce_sum.344 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.345 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %broadcast.879.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1141.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %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} - 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) +%fused_computation.431 (param_0.1531: f32[2048,4], param_1.1712: f32[], param_2.1318: f32[], param_3.880: f32[], param_4.518: f32[2048,4], param_5.462: f32[], param_6.324: f32[4,2048], param_7.198: pred[], param_8.120: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { + %param_0.1531 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.880 = f32[]{:T(128)S(6)} parameter(3) + %mul.2019.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.880), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.198 = pred[]{:T(512)S(6)} parameter(7) + %select_n.280.clone.1 = pred[2048,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.324 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) + %bitcast.691.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.324), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.462 = f32[]{:T(128)} parameter(5) + %div.904.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.462), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.903.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.691.clone.1, %div.904.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.279.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.280.clone.1, %bitcast.691.clone.1, %div.903.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1394.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1009.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1394.clone.1), dimensions={}, metadata={op_name="broadcast.84"} + %mul.2023.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.279.clone.1, %broadcast.1009.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.1398.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.1008.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1398.clone.1), dimensions={}, metadata={op_name="broadcast.83"} + %mul.2022.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.1008.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.967.clone.1 = f32[2048,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.1318 = f32[]{:T(128)S(6)} parameter(2) + %div.900.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.68.clone.1 = f32[2048,4]{0,1:T(4,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.1397.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.1007.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1397.clone.1), dimensions={}, metadata={op_name="broadcast.73"} + %mul.2021.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.68.clone.1, %broadcast.1007.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.518 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) + %constant.1396.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.1006.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1396.clone.1), dimensions={}, metadata={op_name="broadcast.72"} + %mul.2020.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.518, %broadcast.1006.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.966.clone.1 = f32[2048,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.1712 = f32[]{:T(128)S(6)} parameter(1) + %div.899.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1712), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.898.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.966.clone.1, %div.899.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.65.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.898.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1395.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.1004.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1395.clone.1), dimensions={}, metadata={op_name="broadcast.66"} + %add.965.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.65.clone.1, %broadcast.1004.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.451.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.900.clone.1, %add.965.clone.1), metadata={op_name="multiply.52"} + %div.897.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.967.clone.1, %multiply.451.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2018.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1531, %broadcast.1009.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.964.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.897.clone.1, %mul.2018.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2017.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.2019.clone.1, %add.964.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.963.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1531, %mul.2017.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.251 = f32[2048,4]{0,1:T(4,128)} multiply(%add.963.clone.1, %add.963.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1475 = f32[]{:T(128)} constant(0) + %reduce.228 = f32[]{:T(128)} reduce(%square.251, %constant.1475), dimensions={0,1}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.230.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1475), dimensions={0,1}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.159 = (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.228, %add.963.clone.1, %add.966.clone.1, %add.967.clone.1, %reduce.230.clone.1) } %region_11.14 (reduce_sum.192: f32[], reduce_sum.193: f32[]) -> f32[] { - %reduce_sum.193 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} %reduce_sum.192 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.193 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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.442 (param_0.1545: bf16[2048]) -> f32[] { + %param_0.1545 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.1314 = f32[2048]{0:T(1024)} convert(%param_0.1545), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %square.254 = f32[2048]{0:T(1024)} multiply(%convert_element_type.1314, %convert_element_type.1314), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1489 = f32[]{:T(128)} constant(0) + ROOT %reduce.231 = f32[]{:T(128)} reduce(%square.254, %constant.1489), 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[] { - %reduce_sum.402 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_56.61 (reduce_sum.401: f32[], reduce_sum.402: f32[]) -> f32[] { %reduce_sum.401 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.402 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.403 = f32[]{:T(128)} add(%reduce_sum.401, %reduce_sum.402), 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":[]}} } -%region_44.49 (reduce_sum.323: f32[], reduce_sum.324: f32[]) -> f32[] { - %reduce_sum.324 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_41.46 (reduce_sum.323: f32[], reduce_sum.324: f32[]) -> f32[] { %reduce_sum.323 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.324 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %add.1003.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1165.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %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} - 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] { - %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) - %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} - %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) - %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} - %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} - 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} -} - -%region_40.45 (reduce_sum.305: f32[], reduce_sum.309: f32[]) -> f32[] { - %reduce_sum.309 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.443 (param_0.1535: f32[2048], param_1.1716: f32[], param_2.1322: f32[], param_3.884: f32[], param_4.522: f32[2048], param_5.466: f32[], param_6.328: bf16[2048], param_7.202: pred[], param_8.124: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { + %param_0.1535 = f32[2048]{0:T(1024)S(1)} parameter(0) + %param_3.884 = f32[]{:T(128)S(6)} parameter(3) + %mul.2026.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_3.884), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.202 = pred[]{:T(512)S(6)} parameter(7) + %select_n.284.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.328 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) + %convert_element_type.1334.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.328), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_5.466 = f32[]{:T(128)} parameter(5) + %div.912.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.466), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.911.clone.1 = f32[2048]{0:T(1024)} divide(%convert_element_type.1334.clone.1, %div.912.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.283.clone.1 = f32[2048]{0:T(1024)} select(%select_n.284.clone.1, %convert_element_type.1334.clone.1, %div.911.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1400.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.1011.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1400.clone.1), dimensions={}, metadata={op_name="broadcast.88"} + %mul.2032.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.283.clone.1, %broadcast.1011.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.124 = f32[2048]{0:T(1024)S(1)} parameter(8) + %constant.1404.clone.1 = f32[]{:T(128)} constant(0.9) + %mul.2033.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1404.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2031.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.124, %mul.2033.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.973.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2032.clone.1, %mul.2031.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) + %div.908.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.69.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.283.clone.1, %select_n.283.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} + %constant.1403.clone.1 = f32[]{:T(128)} constant(0.05) + %mul.2030.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1403.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.2028.clone.1 = f32[2048]{0:T(1024)} multiply(%integer_pow.69.clone.1, %mul.2030.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.522 = f32[2048]{0:T(1024)S(1)} parameter(4) + %constant.1402.clone.1 = f32[]{:T(128)} constant(0.95) + %mul.2029.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1402.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_4.522, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.972.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_1.1716 = f32[]{:T(128)S(6)} parameter(1) + %div.907.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1716), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.906.clone.1 = f32[2048]{0:T(1024)} divide(%add.972.clone.1, %div.907.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.66.clone.1 = f32[2048]{0:T(1024)} sqrt(%div.906.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1401.clone.1 = f32[]{:T(128)} constant(1e-08) + %add.971.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1401.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %add.970.clone.1 = f32[2048]{0:T(1024)} add(%sqrt.66.clone.1, %add.971.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.452.clone.1 = f32[2048]{0:T(1024)} multiply(%div.908.clone.1, %add.970.clone.1), metadata={op_name="multiply.48"} + %div.905.clone.1 = f32[2048]{0:T(1024)} divide(%add.973.clone.1, %multiply.452.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2025.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1535, %broadcast.1011.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.969.clone.1 = f32[2048]{0:T(1024)} add(%div.905.clone.1, %mul.2025.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2024.clone.1 = f32[2048]{0:T(1024)} multiply(%mul.2026.clone.1, %add.969.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.968.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1535, %mul.2024.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.255 = f32[2048]{0:T(1024)} multiply(%add.968.clone.1, %add.968.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1479 = f32[]{:T(128)} constant(0) + %reduce.232 = f32[]{:T(128)} reduce(%square.255, %constant.1479), dimensions={0}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.233.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1479), dimensions={0}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.162 = (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.232, %add.968.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.233.clone.1) +} + +%fused_computation.448 (param_0.1269: s32[512]) -> s32[1024] { + %constant.1240 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.952 = s32[1024]{0:T(1024)} broadcast(%constant.1240), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %param_0.1269 = s32[512]{0:T(512)S(1)} parameter(0) + %constant.1241 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %pad.57 = s32[1024]{0:T(1024)} pad(%param_0.1269, %constant.1241), padding=0_512, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %constant.1239 = s32[] constant(151935), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.951 = s32[1024]{0:T(1024)} broadcast(%constant.1239), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %clamp.3 = s32[1024]{0:T(1024)S(1)} clamp(%broadcast.952, %pad.57, %broadcast.951), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} +} + +%region_37.42 (reduce_sum.305: f32[], reduce_sum.309: f32[]) -> f32[] { %reduce_sum.305 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.309 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.310 = f32[]{:T(128)} add(%reduce_sum.305, %reduce_sum.309), 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":[]}} } -%region_37.42 (reduce_sum.290: f32[], reduce_sum.291: f32[]) -> f32[] { - %reduce_sum.291 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_34.39 (reduce_sum.290: f32[], reduce_sum.291: f32[]) -> f32[] { %reduce_sum.290 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.291 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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} - ROOT %tuple.170 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.214, %reduce.215.clone.1) +%fused_computation.453 (param_0.1536: bf16[4,128], param_1.1717: bf16[4,128]) -> (f32[], f32[]) { + %param_0.1536 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %convert.107 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1536), metadata={op_name="jit(train_step)/transpose(jvp())/while" stack_frame_id=0} + %bitcast.609 = f32[128,4]{0,1:T(4,128)} bitcast(%convert.107), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.258 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.609, %bitcast.609), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1480 = f32[]{:T(128)} constant(0) + %reduce.234 = f32[]{:T(128)} reduce(%square.258, %constant.1480), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %param_1.1717 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) + %convert.111.clone.1 = f32[4,128]{1,0:T(4,128)} convert(%param_1.1717), metadata={op_name="jit(train_step)/transpose(jvp())/while" stack_frame_id=0} + %bitcast.613.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%convert.111.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %square.261.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.613.clone.1, %bitcast.613.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %reduce.235.clone.1 = f32[]{:T(128)} reduce(%square.261.clone.1, %constant.1480), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.175 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.234, %reduce.235.clone.1) } -%region_72.77 (reduce_sum.470: f32[], reduce_sum.471: f32[]) -> f32[] { - %reduce_sum.471 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.470 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.472 = f32[]{:T(128)} add(%reduce_sum.470, %reduce_sum.471), 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":[]}} -} - -%region_58.63 (reduce_sum.395: f32[], reduce_sum.396: f32[]) -> f32[] { - %reduce_sum.396 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.395 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - 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]) { - %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} - %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} - %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} - 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) -} - -%region_69.74 (reduce_sum.452: f32[], reduce_sum.456: f32[]) -> f32[] { - %reduce_sum.456 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_66.71 (reduce_sum.452: f32[], reduce_sum.456: f32[]) -> f32[] { %reduce_sum.452 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.456 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.457 = f32[]{:T(128)} add(%reduce_sum.452, %reduce_sum.456), 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":[]}} } -%region_54.59 (reduce_sum.374: f32[], reduce_sum.375: f32[]) -> f32[] { - %reduce_sum.375 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_51.56 (reduce_sum.374: f32[], reduce_sum.375: f32[]) -> f32[] { %reduce_sum.374 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.375 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %broadcast.851.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1105.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %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} - 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) -} - -%region_66.71 (reduce_sum.437: f32[], reduce_sum.438: f32[]) -> f32[] { - %reduce_sum.438 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%fused_computation.456 (param_0.1525: f32[128,4], param_1.1706: f32[], param_2.1312: f32[], param_3.874: f32[], param_4.512: f32[128,4], param_5.456: f32[], param_6.318: bf16[4,128], param_7.192: pred[], param_8.114: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1525 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.874 = f32[]{:T(128)S(6)} parameter(3) + %mul.1998.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.874), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.192 = pred[]{:T(512)S(6)} parameter(7) + %select_n.268.clone.1 = pred[128,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.318 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(6) + %convert.169.clone.1 = f32[4,128]{1,0:T(4,128)} convert(%param_6.318), metadata={op_name="jit(train_step)/transpose(jvp())/while" stack_frame_id=0} + %bitcast.685.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%convert.169.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.456 = f32[]{:T(128)} parameter(5) + %div.880.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.456), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.879.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.685.clone.1, %div.880.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.267.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.268.clone.1, %bitcast.685.clone.1, %div.879.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1376.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.991.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1376.clone.1), dimensions={}, metadata={op_name="broadcast.80"} + %mul.2002.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.267.clone.1, %broadcast.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.114 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %constant.1380.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.990.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1380.clone.1), dimensions={}, metadata={op_name="broadcast.79"} + %mul.2001.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.114, %broadcast.990.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.952.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2002.clone.1, %mul.2001.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1312 = f32[]{:T(128)S(6)} parameter(2) + %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1312), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.65.clone.1 = f32[128,4]{0,1:T(4,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.1379.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.989.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1379.clone.1), dimensions={}, metadata={op_name="broadcast.69"} + %mul.2000.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.65.clone.1, %broadcast.989.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.512 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %constant.1378.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.988.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1378.clone.1), dimensions={}, metadata={op_name="broadcast.68"} + %mul.1999.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.512, %broadcast.988.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.2000.clone.1, %mul.1999.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1706 = f32[]{:T(128)S(6)} parameter(1) + %div.875.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1706), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.874.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.951.clone.1, %div.875.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.62.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.874.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1377.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.986.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1377.clone.1), dimensions={}, metadata={op_name="broadcast.64"} + %add.950.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.62.clone.1, %broadcast.986.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.448.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.876.clone.1, %add.950.clone.1), metadata={op_name="multiply.58"} + %div.873.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.952.clone.1, %multiply.448.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.1997.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1525, %broadcast.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.949.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.873.clone.1, %mul.1997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.1996.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1998.clone.1, %add.949.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)S(1)} add(%param_0.1525, %mul.1996.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.262 = f32[128,4]{0,1:T(4,128)} multiply(%add.948.clone.1, %add.948.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1469 = f32[]{:T(128)} constant(0) + %reduce.236 = f32[]{:T(128)} reduce(%square.262, %constant.1469), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.238.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1469), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.165 = (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.236, %add.948.clone.1, %add.951.clone.1, %add.952.clone.1, %reduce.238.clone.1) +} + +%region_63.68 (reduce_sum.437: f32[], reduce_sum.438: f32[]) -> f32[] { %reduce_sum.437 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.438 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} ROOT %reduce_sum.442 = f32[]{:T(128)} add(%reduce_sum.437, %reduce_sum.438), 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":[]}} } -%region_51.56 (reduce_sum.359: f32[], reduce_sum.360: f32[]) -> f32[] { - %reduce_sum.360 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} +%region_48.53 (reduce_sum.359: f32[], reduce_sum.360: f32[]) -> f32[] { %reduce_sum.359 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.360 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} 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) - %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} - %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) - %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} - %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) - %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} - %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) - %broadcast.861.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1123.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %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} - 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) - %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} - %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} - %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} +%fused_computation.457 (param_0.1528: f32[128,4], param_1.1709: f32[], param_2.1315: f32[], param_3.877: f32[], param_4.515: f32[128,4], param_5.459: f32[], param_6.321: bf16[4,128], param_7.195: pred[], param_8.117: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { + %param_0.1528 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) + %param_3.877 = f32[]{:T(128)S(6)} parameter(3) + %mul.2005.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.877), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_7.195 = pred[]{:T(512)S(6)} parameter(7) + %select_n.272.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.195), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %param_6.321 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(6) + %convert.171.clone.1 = f32[4,128]{1,0:T(4,128)} convert(%param_6.321), metadata={op_name="jit(train_step)/transpose(jvp())/while" stack_frame_id=0} + %bitcast.687.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%convert.171.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/transpose" stack_frame_id=0} + %param_5.459 = f32[]{:T(128)} parameter(5) + %div.888.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.459), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.887.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.687.clone.1, %div.888.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %select_n.271.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.272.clone.1, %bitcast.687.clone.1, %div.887.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + %constant.1382.clone.1 = f32[]{:T(128)} constant(0.1) + %broadcast.997.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1382.clone.1), dimensions={}, metadata={op_name="broadcast.80"} + %mul.2009.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.271.clone.1, %broadcast.997.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_8.117 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) + %constant.1386.clone.1 = f32[]{:T(128)} constant(0.9) + %broadcast.996.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1386.clone.1), dimensions={}, metadata={op_name="broadcast.79"} + %mul.2008.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.117, %broadcast.996.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.957.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2009.clone.1, %mul.2008.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) + %div.884.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %integer_pow.66.clone.1 = f32[128,4]{0,1:T(4,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.1385.clone.1 = f32[]{:T(128)} constant(0.05) + %broadcast.995.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1385.clone.1), dimensions={}, metadata={op_name="broadcast.69"} + %mul.2007.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.995.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_4.515 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) + %constant.1384.clone.1 = f32[]{:T(128)} constant(0.95) + %broadcast.994.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1384.clone.1), dimensions={}, metadata={op_name="broadcast.68"} + %mul.2006.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.515, %broadcast.994.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.956.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.2007.clone.1, %mul.2006.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %param_1.1709 = f32[]{:T(128)S(6)} parameter(1) + %div.883.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1709), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.882.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.956.clone.1, %div.883.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sqrt.63.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.882.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} + %constant.1383.clone.1 = f32[]{:T(128)} constant(1e-08) + %broadcast.992.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1383.clone.1), dimensions={}, metadata={op_name="broadcast.64"} + %add.955.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.63.clone.1, %broadcast.992.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %multiply.449.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.884.clone.1, %add.955.clone.1), metadata={op_name="multiply.55"} + %div.881.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.957.clone.1, %multiply.449.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %mul.2004.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1528, %broadcast.997.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.954.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.881.clone.1, %mul.2004.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %mul.2003.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.2005.clone.1, %add.954.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %add.953.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1528, %mul.2003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %square.263 = f32[128,4]{0,1:T(4,128)} multiply(%add.953.clone.1, %add.953.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.1472 = f32[]{:T(128)} constant(0) + %reduce.237 = f32[]{:T(128)} reduce(%square.263, %constant.1472), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + %reduce.239.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1472), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} + ROOT %tuple.166 = (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.237, %add.953.clone.1, %add.956.clone.1, %add.957.clone.1, %reduce.239.clone.1) +} + +%fused_computation.467 (param_0.1364: f32[128]) -> f32[128] { + %param_0.1364 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1314 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.972 = f32[128]{0:T(128)} broadcast(%constant.1314), dimensions={}, metadata={op_name="broadcast.418"} + %div.784 = f32[128]{0:T(128)} multiply(%param_0.1364, %broadcast.972), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.1312 = f32[]{:T(128)} constant(1e-06) + %broadcast.970 = f32[128]{0:T(128)} broadcast(%constant.1312), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.940 = f32[128]{0:T(128)} add(%div.784, %broadcast.970), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %bitcast.624 = f32[1,128]{1,0:T(1,128)} bitcast(%add.940), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %rsqrt.195 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.624), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} + ROOT %bitcast.614 = f32[128]{0:T(128)} bitcast(%rsqrt.195), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} } %region_0.1 (reduce_sum.137: s32[], reduce_sum.138: s32[]) -> s32[] { - %reduce_sum.138 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} %reduce_sum.137 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.138 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} 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} - %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) - %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} - %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.472 (param_0.1369: s32[1,128]) -> s32[] { + %param_0.1369 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.1237 = s32[]{:T(128)} constant(0) + %broadcast.983 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1237), dimensions={}, metadata={op_name="broadcast.99"} + %ne.10 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1369, %broadcast.983), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %convert_element_type.1323 = s32[1,128]{1,0:T(1,128)} convert(%ne.10), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + ROOT %reduce.240 = s32[]{:T(128)} reduce(%convert_element_type.1323, %constant.1237), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.431 () -> f32[64] { - %constant.1049 = f32[]{:T(128)} constant(1e+06) - %broadcast.840 = f32[64]{0:T(128)} broadcast(%constant.1049), dimensions={}, metadata={op_name="broadcast.390"} - %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} - %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} - %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) +%region_69.74 (reduce_sum.470: f32[], reduce_sum.471: f32[]) -> f32[] { + %reduce_sum.470 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.471 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.472 = f32[]{:T(128)} add(%reduce_sum.470, %reduce_sum.471), 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.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) +%region_55.60 (reduce_sum.395: f32[], reduce_sum.396: f32[]) -> f32[] { + %reduce_sum.395 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.396 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + 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.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.475 (param_0.1547: s32[1,128], param_1.1722: bf16[128], param_2.1325: f32[128], param_3.885: f32[128], param_4.523: f32[]) -> (f32[], f32[], f32[128]) { + %param_0.1547 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.1230 = s32[]{:T(128)} constant(0) + %broadcast.981 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1230), dimensions={}, metadata={op_name="broadcast.99"} + %ne.16 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1547, %broadcast.981), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %param_2.1325 = f32[128]{0:T(128)S(1)} parameter(2) + %log.18 = f32[128]{0:T(128)} log(%param_2.1325), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_1.1722 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %reduce_max.16 = f32[128]{0:T(128)} convert(%param_1.1722), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %add.943 = f32[128]{0:T(128)} add(%log.18, %reduce_max.16), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %square.266 = f32[128]{0:T(128)} multiply(%add.943, %add.943), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.1491 = f32[]{:T(128)} constant(0) + %broadcast.966 = f32[128]{0:T(128)} broadcast(%constant.1491), dimensions={}, metadata={op_name="broadcast.120"} + %mul.1981 = f32[128]{0:T(128)} multiply(%square.266, %broadcast.966), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %bitcast.617 = f32[1,128]{1,0:T(1,128)} bitcast(%mul.1981), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %broadcast.977 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1491), dimensions={}, metadata={op_name="broadcast.120"} + %mul.1967 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %bitcast.617, %broadcast.977), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.242 = f32[]{:T(128)} reduce(%mul.1967, %constant.1491), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_3.885 = f32[128]{0:T(128)S(1)} parameter(3) + %neg.113.clone.1 = f32[128]{0:T(128)} negate(%param_3.885), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %bitcast.621.clone.1 = f32[1,128]{1,0:T(1,128)} bitcast(%neg.113.clone.1), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.933.clone.1 = f32[1,128]{1,0:T(1,128)} add(%bitcast.621.clone.1, %bitcast.617), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1965.clone.1 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %add.933.clone.1, %broadcast.977), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.241.clone.1 = f32[]{:T(128)} reduce(%mul.1965.clone.1, %constant.1491), dimensions={0,1}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_4.523 = f32[]{:T(128)S(6)} parameter(4) + %broadcast_in_dim.299.clone.1 = f32[1,128]{1,0:T(1,128)} broadcast(%param_4.523), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} + %mul.1963.clone.1 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %broadcast_in_dim.299.clone.1, %broadcast.977), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %bitcast.615.clone.1 = f32[128]{0:T(128)S(1)} bitcast(%mul.1963.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %tuple.168 = (f32[]{:T(128)}, f32[]{:T(128)}, f32[128]{0:T(128)S(1)}) tuple(%reduce.242, %reduce.241.clone.1, %bitcast.615.clone.1) +} + +%fused_computation.478 (param_0.1336: s32[1,128]) -> s32[1,1,128] { + %param_0.1336 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.1238 = s32[]{:T(128)} constant(0) + %broadcast.982 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1238), dimensions={}, metadata={op_name="broadcast.99"} + %lt.32 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1336, %broadcast.982), direction=LT, metadata={op_name="jit(train_step)/jvp()/lt" stack_frame_id=0} + %constant.1319 = s32[]{:T(128)} constant(151936) + %add.941 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1319), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.930 = s32[1,128]{1,0:T(1,128)} add(%param_0.1336, %add.941), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %select_n.180 = s32[1,128]{1,0:T(1,128)} select(%lt.32, %add.930, %param_0.1336), metadata={op_name="jit(train_step)/jvp()/select_n" stack_frame_id=0} + ROOT %bitcast.616 = s32[1,1,128]{2,1,0:T(1,128)S(1)} bitcast(%select_n.180) +} + +%fused_computation.484 (param_0.1366: f32[128], param_1.1477: f32[128]) -> f32[128] { + %param_0.1366 = f32[128]{0:T(128)S(1)} parameter(0) + %bitcast.631 = f32[1,128]{1,0:T(1,128)} bitcast(%param_0.1366), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0} + %param_1.1477 = f32[128]{0:T(128)S(1)} parameter(1) + %constant.1313 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.971 = f32[128]{0:T(128)} broadcast(%constant.1313), dimensions={}, metadata={op_name="broadcast.418"} + %div.782 = f32[128]{0:T(128)} multiply(%param_1.1477, %broadcast.971), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.1311 = f32[]{:T(128)} constant(1e-06) + %broadcast.969 = f32[128]{0:T(128)} broadcast(%constant.1311), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.938 = f32[128]{0:T(128)} add(%div.782, %broadcast.969), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %bitcast.630 = f32[1,128]{1,0:T(1,128)} bitcast(%add.938), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %rsqrt.201 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.630), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} + %div.780 = f32[1,128]{1,0:T(1,128)} divide(%rsqrt.201, %bitcast.630), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.1308 = f32[]{:T(128)} constant(-0.5) + %mul.1985 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1308), dimensions={}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1979 = f32[1,128]{1,0:T(1,128)} multiply(%div.780, %mul.1985), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1978 = f32[1,128]{1,0:T(1,128)} multiply(%bitcast.631, %mul.1979), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.1307 = f32[]{:T(128)} constant(0.0009765625) + %mul.1984 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1307), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1974 = f32[1,128]{1,0:T(1,128)} multiply(%mul.1978, %mul.1984), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %bitcast.625 = f32[128]{0:T(128)S(1)} bitcast(%mul.1974), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +} + +%fused_computation.485 () -> f32[64] { + %constant.1317 = f32[]{:T(128)} constant(1e+06) + %broadcast.975 = f32[64]{0:T(128)} broadcast(%constant.1317), dimensions={}, metadata={op_name="broadcast.409"} + %iota.56 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/iota" stack_frame_id=0} + %constant.1316 = s32[]{:T(128)} constant(2) + %broadcast.974 = s32[64]{0:T(128)} broadcast(%constant.1316), dimensions={}, metadata={op_name="broadcast.410"} + %mul.1975 = s32[64]{0:T(128)} multiply(%iota.56, %broadcast.974), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %convert_element_type.1324 = f32[64]{0:T(128)} convert(%mul.1975), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %constant.1315 = f32[]{:T(128)} constant(0.0078125) + %broadcast.973 = f32[64]{0:T(128)} broadcast(%constant.1315), dimensions={}, metadata={op_name="broadcast.411"} + %div.778 = f32[64]{0:T(128)} multiply(%convert_element_type.1324, %broadcast.973), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.975, %div.778), metadata={op_name="jit(train_step)/pow" stack_frame_id=0} +} + +%fused_computation.486 (param_0.1352: s32[1,128]) -> (f32[1,128,1,1], f32[128]) { + %param_0.1352 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %convert_element_type.1325 = f32[1,128]{1,0:T(1,128)} convert(%param_0.1352), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %bitcast.626 = f32[1,128,1,1]{1,3,2,0:T(1,128)} bitcast(%convert_element_type.1325), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.627.clone.1 = f32[128]{0:T(128)S(1)} bitcast(%convert_element_type.1325), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + ROOT %tuple.170 = (f32[1,128,1,1]{1,3,2,0:T(1,128)}, f32[128]{0:T(128)S(1)}) tuple(%bitcast.626, %bitcast.627.clone.1) +} + +%region_42.47.clone (reduce_sum.639: f32[], reduce_sum.640: f32[]) -> f32[] { + %reduce_sum.639 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.640 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.641 = f32[]{:T(128)} add(%reduce_sum.639, %reduce_sum.640), 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":[]}} +} + +%region_28.33.clone (reduce_sum.615: f32[], reduce_sum.616: f32[]) -> f32[] { + %reduce_sum.615 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.616 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.617 = f32[]{:T(128)} add(%reduce_sum.615, %reduce_sum.616), 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":[]}} +} + +%region_10.13.clone (dot_general.291: bf16[], dot_general.292: bf16[]) -> bf16[] { + %dot_general.291 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + %dot_general.292 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + ROOT %add.610 = bf16[]{:T(256)} add(%dot_general.291, %dot_general.292), metadata={op_name="add.79"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.491 (param_0.1515: f32[2048,4]) -> bf16[4,2048] { + %param_0.1515 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.755 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1515), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.172 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.755) +} + +%fused_computation.492 (param_0.1514: f32[2048,4]) -> bf16[4,2048] { + %param_0.1514 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.754 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1514), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.174 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.754) +} + +%fused_computation.493 (param_0.1516: f32[128,4]) -> bf16[4,128] { + %param_0.1516 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.756 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1516), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.176 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.756) +} + +%fused_computation.494 (param_0.1517: f32[128,4]) -> bf16[4,128] { + %param_0.1517 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.757 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1517), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.178 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.757) +} + +%fused_computation.495 (param_0.1506: f32[128], param_1.1724: bf16[128]) -> (f32[128], f32[128]) { + %param_0.1506 = f32[128]{0:T(128)S(1)} parameter(0) + %log.23 = f32[128]{0:T(128)S(1)} log(%param_0.1506), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_1.1724 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %reduce_max.18.clone.1 = f32[128]{0:T(128)} convert(%param_1.1724), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %add.945.clone.1 = f32[128]{0:T(128)} add(%log.23, %reduce_max.18.clone.1), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %constant.1493 = f32[]{:T(128)} constant(0) + %broadcast.968.clone.1 = f32[128]{0:T(128)} broadcast(%constant.1493), dimensions={}, metadata={op_name="broadcast.120"} + %mul.1973.clone.1 = f32[128]{0:T(128)} multiply(%add.945.clone.1, %broadcast.968.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.1333.clone.1 = f32[]{:T(128)} constant(1) + %broadcast.965.clone.1 = f32[128]{0:T(128)} broadcast(%constant.1333.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.934.clone.1 = f32[128]{0:T(128)S(1)} add(%mul.1973.clone.1, %broadcast.965.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + ROOT %tuple.169 = (f32[128]{0:T(128)S(1)}, f32[128]{0:T(128)S(1)}) tuple(%log.23, %add.934.clone.1) } -%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) +%region_8.11 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { + %reduce_max.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_max"} + %reduce_max.8 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_max"} + 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.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.412.clone.clone (param_0.1510: bf16[2048], param_1.1694: f32[128], param_2.1284: bf16[1,128,2048]) -> bf16[128,2048] { + %param_0.1510 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.506 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1510), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %param_2.1284 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} parameter(2) + %convert_element_type.1364 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_2.1284), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1694 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2139 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1694), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.2138 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1364, %mul.2139), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.1363 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2138), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %dot_general.505 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.506, %convert_element_type.1363), metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + ROOT %bitcast.749 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%dot_general.505) } -%region_8.11 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { - %reduce_max.8 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_max"} - %reduce_max.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_max"} - 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":[]}} +%bitcast_fusion.3 (bitcast_input.3: bf16[151936,2048]) -> bf16[151936,2048] { + %bitcast_input.3 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.761 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.3) } -%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.496 (param_0.1518: bf16[151936,2048], param_1.1699: bf16[2048], param_2.1301: f32[128], param_3.867: bf16[1,128,2048]) -> (bf16[128], bf16[128,151936]) { + %param_1.1699 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1301 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.867 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} parameter(3) + %fusion.319.clone.1 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_1.1699, %param_2.1301, %param_3.867), kind=kLoop, calls=%fused_computation.412.clone.clone + %param_0.1518 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %fusion.357 = bf16[151936,2048]{1,0:T(8,128)(2,1)} fusion(%param_0.1518), kind=kLoop, calls=%bitcast_fusion.3 + %convolution.115.clone.1 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} convolution(%fusion.319.clone.1, %fusion.357), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/jvp()/dot_general" stack_frame_id=0} + %constant.1462 = bf16[]{:T(256)} constant(-inf) + %reduce.243 = bf16[128]{0:T(256)(128)(2,1)S(1)} reduce(%convolution.115.clone.1, %constant.1462), dimensions={1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + ROOT %tuple.171 = (bf16[128]{0:T(256)(128)(2,1)S(1)}, bf16[128,151936]{1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.243, %convolution.115.clone.1) } -%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} +%convert_element_type.762.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { + %lhs.1 = bf16[] parameter(0) + %rhs.1 = bf16[] parameter(1) + ROOT %add.793 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%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.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) +%convert_element_type.758.reduce_sub_computation (lhs: bf16[], rhs: bf16[]) -> bf16[] { + %lhs = bf16[] parameter(0) + %rhs = bf16[] parameter(1) + ROOT %add.792 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%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.285.clone.clone (param_0.1701: bf16[4,2048], param_1.1822: s32[], param_2.1394: bf16[4,2048]) -> (bf16[2048], bf16[2048]) { + %param_0.1701 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1822 = s32[]{:T(128)S(6)} parameter(1) + %constant.1626 = s32[]{:T(128)} constant(0) + %dynamic_slice.289 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1701, %param_1.1822, %constant.1626), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.1627 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %reduce.263 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.289, %constant.1627), dimensions={0}, to_apply=%convert_element_type.762.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_2.1394 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(2) + %dynamic_slice.274.clone.3 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_2.1394, %param_1.1822, %constant.1626), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} + %reduce.195.clone.3 = bf16[2048]{0:T(1024)(128)(2,1)} reduce(%dynamic_slice.274.clone.3, %constant.1627), dimensions={0}, to_apply=%convert_element_type.758.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %tuple.216 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[2048]{0:T(1024)(128)(2,1)}) tuple(%reduce.263, %reduce.195.clone.3) } -%convert_element_type.767.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { - %rhs.1 = bf16[] parameter(1) - %lhs.1 = bf16[] parameter(0) - 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.130.clone.clone (param_0.1702: bf16[4,512,8,128], param_1.1823: s32[]) -> bf16[1,512,8,128] { + %param_0.1702 = bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1823 = s32[]{:T(128)S(6)} parameter(1) + %constant.1628 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.185 = bf16[1,512,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1702, %param_1.1823, %constant.1628, %constant.1628, %constant.1628), dynamic_slice_sizes={1,512,8,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.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.121.clone.clone (param_0.1718: bf16[4,512,16,128], param_1.1833: s32[]) -> bf16[1,512,16,128] { + %param_0.1718 = bf16[4,512,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1833 = s32[]{:T(128)S(6)} parameter(1) + %constant.1644 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.189 = bf16[1,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1718, %param_1.1833, %constant.1644, %constant.1644, %constant.1644), dynamic_slice_sizes={1,512,16,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} } %region_14.16 (reduce_sum.204: f32[], reduce_sum.205: f32[]) -> f32[] { - %reduce_sum.205 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.204 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - 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} - %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} -} - -%region_15.17 (reduce_sum.207: f32[], reduce_sum.211: f32[]) -> f32[] { - %reduce_sum.211 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.207 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - 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} - %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.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.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.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.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.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.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.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.1563, %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} - %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} -} - -%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.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.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.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.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[] { - %reduce_sum.214 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.213 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - 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} - %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} - 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} - %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.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.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[] { - %reduce_sum.220 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.219 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - 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.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.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.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.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) -} - -%convert_element_type.763.reduce_sub_computation (lhs: bf16[], rhs: bf16[]) -> bf16[] { - %rhs = bf16[] parameter(1) - %lhs = bf16[] parameter(0) - 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.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)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.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.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.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} + %reduce_sum.204 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.205 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.206 = f32[]{:T(128)} add(%reduce_sum.204, %reduce_sum.205), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.172.clone.clone (param_0.1699: bf16[4,1,128,2048], param_1.1821: s32[]) -> f32[128] { + %param_0.1699 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1821 = s32[]{:T(128)S(6)} parameter(1) + %constant.1622 = s32[]{:T(128)} constant(0) + %dynamic-slice.184 = bf16[1,1,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1699, %param_1.1821, %constant.1622, %constant.1622, %constant.1622), dynamic_slice_sizes={1,1,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.902 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.184), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/squeeze" stack_frame_id=0} + %convert_element_type.1459 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.902), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.277 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1459, %convert_element_type.1459), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.1623 = f32[]{:T(128)} constant(0) + ROOT %reduce.262 = f32[128]{0:T(128)S(1)} reduce(%square.277, %constant.1623), dimensions={0,2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} +} + +%fused_computation.299.clone.clone (param_0.1700: f32[128]) -> f32[128] { + %param_0.1700 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1625 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.1073 = f32[128]{0:T(128)} broadcast(%constant.1625), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1009 = f32[128]{0:T(128)} multiply(%param_0.1700, %broadcast.1073), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1624 = f32[]{:T(128)} constant(1e-06) + %broadcast.1072 = f32[128]{0:T(128)} broadcast(%constant.1624), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.1045 = f32[128]{0:T(128)} add(%div.1009, %broadcast.1072), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.904 = f32[1,128]{1,0:T(1,128)} bitcast(%add.1045), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.214 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.904), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.903 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.214), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.279.clone.clone (param_0.1722: f32[128,16]) -> f32[128,16] { + %param_0.1722 = f32[128,16]{0,1:T(8,128)S(1)} parameter(0) + %constant.1647 = f32[]{:T(128)} constant(0.0078125) + %broadcast.1079 = f32[128,16]{0,1:T(8,128)} broadcast(%constant.1647), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1015 = f32[128,16]{0,1:T(8,128)} multiply(%param_0.1722, %broadcast.1079), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1648 = f32[]{:T(128)} constant(1e-06) + %broadcast.1078 = f32[128,16]{0,1:T(8,128)} broadcast(%constant.1648), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %add.1048 = f32[128,16]{0,1:T(8,128)} add(%div.1015, %broadcast.1078), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.923 = f32[1,128,16]{1,2,0:T(8,128)} bitcast(%add.1048), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.216 = f32[1,128,16]{1,2,0:T(8,128)} rsqrt(%bitcast.923), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.922 = f32[128,16]{0,1:T(8,128)S(1)} bitcast(%rsqrt.216), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.305.clone.clone (param_0.1706: bf16[4,128], param_1.1826: s32[], param_2.1397: bf16[4,128]) -> (bf16[128], bf16[128]) { + %param_0.1706 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1826 = s32[]{:T(128)S(6)} parameter(1) + %constant.1630 = s32[]{:T(128)} constant(0) + %dynamic_slice.290 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1706, %param_1.1826, %constant.1630), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} + %bitcast.909 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.290), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_2.1397 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(2) + %dynamic_slice.284.clone.3 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_2.1397, %param_1.1826, %constant.1630), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} + %bitcast.470.clone.3 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.284.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %tuple.217 = (bf16[128]{0:T(256)(128)(2,1)S(1)}, bf16[128]{0:T(256)(128)(2,1)S(1)}) tuple(%bitcast.909, %bitcast.470.clone.3) +} + +%fused_computation.237.clone.1.clone (param_0.1723: f32[128,16], param_1.1836: bf16[128,16,128], param_2.1404: bf16[128]) -> bf16[1,128,16,128] { + %param_2.1404 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.595 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1404), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1836 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1472 = f32[128,16,128]{2,0,1:T(8,128)} convert(%param_1.1836), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.924 = f32[1,128,16,128]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.1472), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_0.1723 = f32[128,16]{0,1:T(8,128)S(1)} parameter(0) + %mul.2312 = f32[1,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1723), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2311 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%bitcast.924, %mul.2312), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1471 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2311), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %dot_general.594 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.595, %convert_element_type.1471), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.198.clone.clone (param_0.1724: bf16[1,128,16,128]) -> (bf16[1,128,16,64], bf16[1,128,16,64]) { + %param_0.1724 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.75 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1724), slice={[0:1], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %neg.127 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%slice.75), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0} + %slice.76 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1724), slice={[0:1], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.224 = (bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.127, %slice.76) +} + +%fused_computation.314.clone.clone () -> f32[64] { + %constant.1637 = f32[]{:T(128)} constant(1e+06) + %closed_call.68 = f32[64]{0:T(128)} broadcast(%constant.1637), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %iota.65 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/iota" stack_frame_id=0} + %constant.1636 = s32[]{:T(128)} constant(2) + %closed_call.67 = s32[64]{0:T(128)} broadcast(%constant.1636), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.2298 = s32[64]{0:T(128)} multiply(%iota.65, %closed_call.67), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1464 = f32[64]{0:T(128)} convert(%mul.2298), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %constant.1638 = f32[]{:T(128)} constant(0.0078125) + %closed_call.66 = f32[64]{0:T(128)} broadcast(%constant.1638), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1011 = f32[64]{0:T(128)} multiply(%convert_element_type.1464, %closed_call.66), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.68, %div.1011), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/pow" stack_frame_id=0} +} + +%fused_computation.270.clone.clone (param_0.1712: f32[128], param_1.1830: f32[64]) -> (bf16[1,128,1,64], bf16[1,128,1,64]) { + %param_0.1712 = f32[128]{0:T(128)S(1)} parameter(0) + %div.1013 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_0.1712), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %param_1.1830 = f32[64]{0:T(128)S(1)} parameter(1) + %div.1014 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_1.1830), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %div.1012 = f32[1,128,1,64]{3,1,2,0:T(8,128)} divide(%div.1013, %div.1014), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %cos.43 = f32[1,128,1,64]{3,1,2,0:T(8,128)} cosine(%div.1012), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/cos" stack_frame_id=0} + %convert_element_type.1465 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %sin.35.clone.3 = f32[1,128,1,64]{3,1,2,0:T(8,128)} sine(%div.1012), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/sin" stack_frame_id=0} + %convert_element_type.1140.clone.3 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %tuple.219 = (bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1465, %convert_element_type.1140.clone.3) +} + +%fused_computation.275.clone.clone (param_0.1714: bf16[1,128,1,64]) -> (bf16[128,128], bf16[128,128]) { + %param_0.1714 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1641 = bf16[]{:T(256)} constant(-inf) + %pad.77 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1714, %constant.1641), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %pad.76 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1714, %constant.1641), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %maximum.57 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.77, %pad.76), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %constant.1642 = bf16[]{:T(256)} constant(0.08838) + %broadcast.1077 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1642), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.2300 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.57, %broadcast.1077), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %bitcast.916 = bf16[128,128]{1,0:T(8,128)(2,1)} bitcast(%mul.2300), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %bitcast.455.clone.3 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.57), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + ROOT %tuple.221 = (bf16[128,128]{1,0:T(8,128)(2,1)}, bf16[128,128]{1,0:T(8,128)(2,1)S(1)}) tuple(%bitcast.916, %bitcast.455.clone.3) +} + +%fused_computation.274.clone.clone (param_0.1713: bf16[1,128,1,64]) -> (bf16[128,128], bf16[128,128]) { + %param_0.1713 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1639 = bf16[]{:T(256)} constant(-inf) + %pad.75 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1713, %constant.1639), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %pad.74 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1713, %constant.1639), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %maximum.56 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.75, %pad.74), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %constant.1640 = bf16[]{:T(256)} constant(0.08838) + %broadcast.1076 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1640), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.2299 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %broadcast.1076), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %bitcast.915 = bf16[128,128]{1,0:T(8,128)(2,1)} bitcast(%mul.2299), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %bitcast.456.clone.3 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.56), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + ROOT %tuple.220 = (bf16[128,128]{1,0:T(8,128)(2,1)}, bf16[128,128]{1,0:T(8,128)(2,1)S(1)}) tuple(%bitcast.915, %bitcast.456.clone.3) +} + +%fused_computation.202.clone.clone (param_0.1725: bf16[1,128,16,64], param_1.1837: bf16[1,128,16,64], param_2.1405: bf16[128,128], param_3.930: bf16[128,128], param_4.551: f32[128,16], param_5.481: bf16[128,16,128], param_6.338: bf16[128]) -> bf16[16,128,128] { + %param_6.338 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.597 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.338), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %param_5.481 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(5) + %convert_element_type.1474 = f32[128,16,128]{2,0,1:T(8,128)} convert(%param_5.481), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.926 = f32[1,128,16,128]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.1474), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_4.551 = f32[128,16]{0,1:T(8,128)S(1)} parameter(4) + %mul.2319 = f32[1,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.551), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2318 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%bitcast.926, %mul.2319), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1473 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2318), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %dot_general.596 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.597, %convert_element_type.1473), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.930 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2317 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.930), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2315 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.596, %mul.2317), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %param_1.1837 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1649 = bf16[]{:T(256)} constant(-inf) + %pad.81 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1837, %constant.1649), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %param_0.1725 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.80 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1725, %constant.1649), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %maximum.59 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.81, %pad.80), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %param_2.1405 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2316 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1405), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2314 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.59, %mul.2316), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %add.1049 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2315, %mul.2314), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %constant.1650 = bf16[]{:T(256)} constant(0.08838) + %closed_call.69 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1650), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.2313 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1049, %closed_call.69), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + ROOT %bitcast.925 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2313), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(jit(_splash_attention))/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.131.clone.clone (param_0.1707: bf16[4,512,8,128], param_1.1827: s32[]) -> bf16[1,512,8,128] { + %param_0.1707 = bf16[4,512,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1827 = s32[]{:T(128)S(6)} parameter(1) + %constant.1631 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.187 = bf16[1,512,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1707, %param_1.1827, %constant.1631, %constant.1631, %constant.1631), dynamic_slice_sizes={1,512,8,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.288.clone.clone (param_0.1711: f32[128,8]) -> f32[128,8] { + %param_0.1711 = f32[128,8]{0,1:T(8,128)S(1)} parameter(0) + %constant.1634 = f32[]{:T(128)} constant(0.0078125) + %broadcast.1075 = f32[128,8]{0,1:T(8,128)} broadcast(%constant.1634), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1010 = f32[128,8]{0,1:T(8,128)} multiply(%param_0.1711, %broadcast.1075), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1635 = f32[]{:T(128)} constant(1e-06) + %broadcast.1074 = f32[128,8]{0,1:T(8,128)} broadcast(%constant.1635), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %add.1046 = f32[128,8]{0,1:T(8,128)} add(%div.1010, %broadcast.1074), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.914 = f32[1,128,8]{1,2,0:T(8,128)} bitcast(%add.1046), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.215 = f32[1,128,8]{1,2,0:T(8,128)} rsqrt(%bitcast.914), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.913 = f32[128,8]{0,1:T(8,128)S(1)} bitcast(%rsqrt.215), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.267.clone.1.clone (param_0.1715: bf16[128], param_1.1831: f32[1,128,8,128], param_2.1400: f32[128,8]) -> bf16[1,128,8,128] { + %param_0.1715 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(0) + %dot_general.589 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_0.1715), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %param_1.1831 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} parameter(1) + %param_2.1400 = f32[128,8]{0,1:T(8,128)S(1)} parameter(2) + %mul.2302 = f32[1,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_2.1400), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2301 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%param_1.1831, %mul.2302), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1466 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2301), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %dot_general.588 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.589, %convert_element_type.1466), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} +} + +%fused_computation.246.clone.clone (param_0.1716: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1716 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.73 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1716), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %neg.126 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%slice.73), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0} + %slice.74 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1716), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.222 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.126, %slice.74) +} + +%fused_computation.249.clone.clone (param_0.1717: bf16[1,128,8,64], param_1.1832: bf16[1,128,8,64], param_2.1401: bf16[128,128], param_3.927: bf16[128,128], param_4.549: f32[1,128,8,128], param_5.480: f32[128,8], param_6.337: bf16[128]) -> bf16[8,128,128] { + %param_6.337 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.591 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.337), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %param_4.549 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} parameter(4) + %param_5.480 = f32[128,8]{0,1:T(8,128)S(1)} parameter(5) + %mul.2308 = f32[1,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_5.480), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2307 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%param_4.549, %mul.2308), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1467 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2307), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %dot_general.590 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.591, %convert_element_type.1467), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.927 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2306 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.927), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2304 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.590, %mul.2306), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %param_1.1832 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1643 = bf16[]{:T(256)} constant(-inf) + %pad.79 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1832, %constant.1643), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %param_0.1717 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.78 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1717, %constant.1643), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %maximum.58 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.79, %pad.78), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %param_2.1401 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2305 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1401), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2303 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.58, %mul.2305), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %add.1047 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2304, %mul.2303), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + ROOT %bitcast.917 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1047), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(jit(_splash_attention))/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.488.clone.clone (param_0.1754: f32[1,16,128,128]) -> (f32[1,16,128,1], f32[16,128]) { + %param_0.1754 = f32[1,16,128,128]{2,1,3,0:T(8,128)S(1)} parameter(0) + %slice.77 = f32[1,16,128,1]{2,1,3,0:T(8,128)} slice(%param_0.1754), slice={[0:1], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} + %bitcast.489.clone.3 = f32[16,128]{1,0:T(8,128)S(1)} bitcast(%slice.77), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/broadcast_in_dim" stack_frame_id=0} + ROOT %tuple.229 = (f32[1,16,128,1]{2,1,3,0:T(8,128)}, f32[16,128]{1,0:T(8,128)S(1)}) tuple(%slice.77, %bitcast.489.clone.3) +} + +%fused_computation.120.clone.clone (param_0.1726: bf16[4,16,128,512], param_1.1838: s32[]) -> bf16[1,16,128,512] { + %param_0.1726 = bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1838 = s32[]{:T(128)S(6)} parameter(1) + %constant.1651 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.191 = bf16[1,16,128,512]{3,2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1726, %param_1.1838, %constant.1651, %constant.1651, %constant.1651), dynamic_slice_sizes={1,16,128,512}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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.98.clone.clone (param_0.1731: bf16[4,512,6144], param_1.1840: s32[]) -> bf16[1,512,6144] { + %param_0.1731 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1840 = s32[]{:T(128)S(6)} parameter(1) + %constant.1655 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.192 = bf16[1,512,6144]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1731, %param_1.1840, %constant.1655, %constant.1655), dynamic_slice_sizes={1,512,6144}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} +} + +%fused_computation.99.clone.clone (param_0.1735: bf16[4,6144,512], param_1.1843: s32[]) -> bf16[1,6144,512] { + %param_0.1735 = bf16[4,6144,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1843 = s32[]{:T(128)S(6)} parameter(1) + %constant.1656 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.193 = bf16[1,6144,512]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1735, %param_1.1843, %constant.1656, %constant.1656), dynamic_slice_sizes={1,6144,512}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} +} + +%fused_computation.100.clone.clone (param_0.1739: bf16[4,512,6144], param_1.1845: s32[]) -> bf16[1,512,6144] { + %param_0.1739 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1845 = s32[]{:T(128)S(6)} parameter(1) + %constant.1657 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.194 = bf16[1,512,6144]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1739, %param_1.1845, %constant.1657, %constant.1657), dynamic_slice_sizes={1,512,6144}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/dynamic_slice" stack_frame_id=0}, 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":[]} +} + +%fused_computation.297.clone.clone (param_0.1730: f32[128]) -> f32[128] { + %param_0.1730 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1654 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.1081 = f32[128]{0:T(128)} broadcast(%constant.1654), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1016 = f32[128]{0:T(128)} multiply(%param_0.1730, %broadcast.1081), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1653 = f32[]{:T(128)} constant(1e-06) + %broadcast.1080 = f32[128]{0:T(128)} broadcast(%constant.1653), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.1050 = f32[128]{0:T(128)} add(%div.1016, %broadcast.1080), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.930 = f32[1,128]{1,0:T(1,128)} bitcast(%add.1050), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.217 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.930), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.929 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.217), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.313.clone.clone (param_0.1749: f32[128], param_1.1852: f32[128]) -> f32[128] { + %param_0.1749 = f32[128]{0:T(128)S(1)} parameter(0) + %bitcast.946 = f32[1,128]{1,0:T(1,128)} bitcast(%param_0.1749), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum" stack_frame_id=0} + %param_1.1852 = f32[128]{0:T(128)S(1)} parameter(1) + %constant.1664 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.1083 = f32[128]{0:T(128)} broadcast(%constant.1664), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1020 = f32[128]{0:T(128)} multiply(%param_1.1852, %broadcast.1083), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1663 = f32[]{:T(128)} constant(1e-06) + %broadcast.1082 = f32[128]{0:T(128)} broadcast(%constant.1663), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.1053 = f32[128]{0:T(128)} add(%div.1020, %broadcast.1082), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.945 = f32[1,128]{1,0:T(1,128)} bitcast(%add.1053), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.218 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.945), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + %div.1019 = f32[1,128]{1,0:T(1,128)} divide(%rsqrt.218, %bitcast.945), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1662 = f32[]{:T(128)} constant(-0.5) + %closed_call.70 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1662), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.2334 = f32[1,128]{1,0:T(1,128)} multiply(%div.1019, %closed_call.70), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2333 = f32[1,128]{1,0:T(1,128)} multiply(%bitcast.946, %mul.2334), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.1661 = f32[]{:T(128)} constant(0.0009765625) + %mul.2335 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.1661), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.2332 = f32[1,128]{1,0:T(1,128)} multiply(%mul.2333, %mul.2335), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + ROOT %bitcast.944 = f32[128]{0:T(128)S(1)} bitcast(%mul.2332), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%region_23.27 (dot_general.209: bf16[], dot_general.210: bf16[]) -> bf16[] { + %dot_general.209 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + %dot_general.210 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + ROOT %add.418 = bf16[]{:T(256)} add(%dot_general.209, %dot_general.210), metadata={op_name="add.70"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.230.clone.clone (param_0.1750: bf16[1,128,2048], param_1.1853: f32[128], param_2.1414: bf16[1,128,2048], param_3.935: bf16[1,128,2048], param_4.553: f32[128], param_5.483: bf16[2048]) -> (bf16[2048], bf16[1,128,2048]) { + %param_0.1750 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_2.1414 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1483 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_2.1414), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1853 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2337 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1853), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2336 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1483, %mul.2337), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1482 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2336), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %multiply.470 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1750, %convert_element_type.1482), metadata={op_name="multiply.352"} + %constant.1665 = bf16[]{:T(256)} constant(0) + %reduce.268 = bf16[2048]{0:T(1024)(128)(2,1)} reduce(%multiply.470, %constant.1665), dimensions={0,1}, to_apply=%region_23.27, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.935 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_5.483 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %dot_general.448.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.483), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.364.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1750, %dot_general.448.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1115.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.364.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %mul.1579.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1115.clone.3, %mul.2337), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %param_4.553 = f32[128]{0:T(128)S(1)} parameter(4) + %mul.1598.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_4.553), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.1578.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1483, %mul.1598.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %add_any.171.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} add(%mul.1579.clone.3, %mul.1578.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %convert_element_type.1113.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.171.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %add_any.169.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.935, %convert_element_type.1113.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + ROOT %tuple.227 = (bf16[2048]{0:T(1024)(128)(2,1)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.268, %add_any.169.clone.3) +} + +%fused_computation.255.clone.clone (param_0.1777: bf16[8,128,128], param_1.1871: bf16[128,128]) -> bf16[1,128,8,128] { + %param_0.1777 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.973 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1777), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/splash_mha_dkv_segmented_no_residuals/splash_mha_dkv_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %param_1.1871 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(1) + %broadcast.1087 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_1.1871), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2359 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} multiply(%bitcast.973, %broadcast.1087), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + ROOT %bitcast.972 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%mul.2359), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%fused_computation.490.clone.clone (param_0.1778: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1778 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.79 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1778), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %slice.52.clone.3 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1778), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %neg.100.clone.3 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%slice.52.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/neg" stack_frame_id=0} + ROOT %tuple.234 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%slice.79, %neg.100.clone.3) +} + +%region_25.29 (dot_general.213: bf16[], dot_general.214: bf16[]) -> bf16[] { + %dot_general.213 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + %dot_general.214 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + ROOT %add.420 = bf16[]{:T(256)} add(%dot_general.213, %dot_general.214), metadata={op_name="add.72"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.256.clone.clone (param_0.1779: f32[1,128,8,128], param_1.1872: f32[128,8], param_2.1428: bf16[128,128], param_3.943: bf16[8,128,128], param_4.559: bf16[1,128,8,64], param_5.486: bf16[1,128,8,64]) -> (bf16[128], bf16[1,128,8,128]) { + %param_5.486 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) + %constant.1375.clone.3 = bf16[]{:T(256)} constant(-inf) + %pad.61.clone.3 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_5.486, %constant.1375.clone.3), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %param_4.559 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(4) + %pad.60.clone.3 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_4.559, %constant.1375.clone.3), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %maximum.49.clone.3 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.61.clone.3, %pad.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %param_3.943 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %bitcast.683.clone.3 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_3.943), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(jit(_splash_attention))/splash_mha_dkv_segmented_no_residuals/splash_mha_dkv_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %param_2.1428 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %broadcast.985.clone.3 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_2.1428), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1995.clone.3 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} multiply(%bitcast.683.clone.3, %broadcast.985.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %bitcast.682.clone.3 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%mul.1995.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %add_any.195.clone.3 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} add(%maximum.49.clone.3, %bitcast.682.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %param_0.1779 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} parameter(0) + %param_1.1872 = f32[128,8]{0,1:T(8,128)S(1)} parameter(1) + %mul.2361 = f32[1,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_1.1872), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2360 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%param_0.1779, %mul.2361), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1497 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2360), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %multiply.475 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add_any.195.clone.3, %convert_element_type.1497), metadata={op_name="multiply.355"} + %constant.1683 = bf16[]{:T(256)} constant(0) + %reduce.271 = bf16[128]{0:T(256)(128)(2,1)} reduce(%multiply.475, %constant.1683), dimensions={0,1,2}, to_apply=%region_25.29, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + ROOT %tuple.235 = (bf16[128]{0:T(256)(128)(2,1)}, bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.271, %add_any.195.clone.3) +} + +%region_20.24 (reduce_sum.228: f32[], reduce_sum.232: f32[]) -> f32[] { + %reduce_sum.228 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + %reduce_sum.232 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + ROOT %reduce_sum.233 = f32[]{:T(128)} add(%reduce_sum.228, %reduce_sum.232), metadata={op_name="checkpoint/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.193.clone.clone (param_0.1780: f32[1,128,8,128], param_1.1873: bf16[1,128,8,128], param_2.1429: bf16[128]) -> f32[128,8] { + %param_0.1780 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} parameter(0) + %param_1.1873 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1429 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.616 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1429), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %dot_general.615 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_1.1873, %dot_general.616), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + %convert_element_type.1498 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%dot_general.615), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %mul.2362 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%param_0.1780, %convert_element_type.1498), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.1684 = f32[]{:T(128)} constant(0) + ROOT %reduce.272 = f32[128,8]{0,1:T(8,128)S(1)} reduce(%mul.2362, %constant.1684), dimensions={0,3}, to_apply=%region_20.24, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum" stack_frame_id=0} +} + +%fused_computation.292.clone.clone (param_0.1781: f32[128,8], param_1.1874: f32[128,8]) -> f32[128,8] { + %param_0.1781 = f32[128,8]{0,1:T(8,128)S(1)} parameter(0)