HloModule jit_program, is_scheduled=true, entry_computation_layout={(f32[256,32]{1,0}, s32[256]{0})->f32[256]{0}}, allow_spmd_sharding_propagation_to_parameters={true,true}, allow_spmd_sharding_propagation_to_output={true}

FileNames
1 "/home/marielle/projects/xtrax/scripts/prof_stage1_onehot_micro.py"

FunctionNames
1 "<module>"
2 "main"
3 "_build_program.<locals>.program"
4 "_kernel"

FileLocations
1 {file_name_id=1 function_name_id=1 line=256 end_line=256 column=21 end_column=27}
2 {file_name_id=1 function_name_id=2 line=169 end_line=169 column=15 end_column=46}
3 {file_name_id=1 function_name_id=3 line=90 end_line=90 column=16 end_column=34}
4 {file_name_id=1 function_name_id=4 line=73 end_line=73 column=11 end_column=48}
5 {file_name_id=1 function_name_id=3 line=92 end_line=92 column=16 end_column=60}
6 {file_name_id=1 function_name_id=4 line=73 end_line=73 column=13 end_column=28}
7 {file_name_id=1 function_name_id=4 line=73 end_line=73 column=12 end_column=34}
8 {file_name_id=1 function_name_id=3 line=92 end_line=92 column=24 end_column=56}
9 {file_name_id=1 function_name_id=3 line=93 end_line=93 column=15 end_column=20}

StackFrames
1 {file_location_id=1 parent_frame_id=1}
2 {file_location_id=2 parent_frame_id=2}
3 {file_location_id=3 parent_frame_id=3}
4 {file_location_id=4 parent_frame_id=4}
5 {file_location_id=5 parent_frame_id=3}
6 {file_location_id=4 parent_frame_id=6}
7 {file_location_id=6 parent_frame_id=4}
8 {file_location_id=7 parent_frame_id=4}
9 {file_location_id=8 parent_frame_id=3}
10 {file_location_id=6 parent_frame_id=6}
11 {file_location_id=7 parent_frame_id=6}
12 {file_location_id=9 parent_frame_id=3}


%region_1.3 (reduce_sum.10: f32[], reduce_sum.11: f32[]) -> f32[] {
  %reduce_sum.10 = f32[] parameter(0), metadata={op_name="reduce_sum"}
  %reduce_sum.11 = f32[] parameter(1), metadata={op_name="reduce_sum"}
  ROOT %reduce_sum.12 = f32[] add(%reduce_sum.10, %reduce_sum.11), metadata={op_name="jit(program)/onehot_onthefly/reduce_sum" stack_frame_id=6}
}

%fused_computation (param_0.1: f32[], param_1.2: f32[256,32], param_2: f32[32,16]) -> f32[256] {
  %param_1.2 = f32[256,32]{1,0} parameter(1)
  %param_2 = f32[32,16]{1,0} parameter(2)
  %dot_general.0 = f32[256,16]{1,0} dot(%param_1.2, %param_2), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(program)/onehot_onthefly/dot_general" stack_frame_id=10}
  %integer_pow.0 = f32[256,16]{1,0} multiply(%dot_general.0, %dot_general.0), metadata={op_name="jit(program)/onehot_onthefly/integer_pow" stack_frame_id=11}
  %param_0.1 = f32[] parameter(0)
  ROOT %reduce_sum.0 = f32[256]{0} reduce(%integer_pow.0, %param_0.1), dimensions={1}, to_apply=%region_1.3, metadata={op_name="jit(program)/onehot_onthefly/reduce_sum" stack_frame_id=6}
}

%region_0.1 (reduce_sum.3: f32[], reduce_sum.4: f32[]) -> f32[] {
  %reduce_sum.3 = f32[] parameter(0), metadata={op_name="reduce_sum"}
  %reduce_sum.4 = f32[] parameter(1), metadata={op_name="reduce_sum"}
  ROOT %reduce_sum.5 = f32[] add(%reduce_sum.3, %reduce_sum.4), metadata={op_name="jit(program)/onehot_materialized/reduce_sum" stack_frame_id=4}
}

%fused_computation.1 (param_0.3: f32[], param_1.5: f32[256,32], param_2.1: f32[32,16]) -> f32[256] {
  %param_1.5 = f32[256,32]{1,0} parameter(1)
  %param_2.1 = f32[32,16]{1,0} parameter(2)
  %dot_general.1 = f32[256,16]{1,0} dot(%param_1.5, %param_2.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_name="jit(program)/onehot_materialized/dot_general" stack_frame_id=7}
  %integer_pow.1 = f32[256,16]{1,0} multiply(%dot_general.1, %dot_general.1), metadata={op_name="jit(program)/onehot_materialized/integer_pow" stack_frame_id=8}
  %param_0.3 = f32[] parameter(0)
  ROOT %reduce_sum.1 = f32[256]{0} reduce(%integer_pow.1, %param_0.3), dimensions={1}, to_apply=%region_0.1, metadata={op_name="jit(program)/onehot_materialized/reduce_sum" stack_frame_id=4}
}

%fused_computation.2 (param_0.6: s32[256]) -> f32[256,32] {
  %param_0.6 = s32[256]{0} parameter(0)
  %eq.15 = s32[256,32]{1,0} broadcast(%param_0.6), dimensions={0}, metadata={op_name="jit(program)/onehot_onthefly/jit(_one_hot)/eq" stack_frame_id=9}
  %iota.7 = s32[256,32]{1,0} iota(), iota_dimension=1, metadata={op_name="jit(program)/onehot_onthefly/jit(_one_hot)/eq" stack_frame_id=9}
  %eq.14 = pred[256,32]{1,0} compare(%eq.15, %iota.7), direction=EQ, metadata={op_name="jit(program)/onehot_onthefly/jit(_one_hot)/eq" stack_frame_id=9}
  ROOT %convert_element_type.2 = f32[256,32]{1,0} convert(%eq.14), metadata={op_name="jit(program)/onehot_onthefly/jit(_one_hot)/convert_element_type" stack_frame_id=9}
}

%wrapped_add_computation (param_0.7: f32[256], param_1.7: f32[256]) -> f32[256] {
  %param_0.7 = f32[256]{0} parameter(0)
  %param_1.7 = f32[256]{0} parameter(1)
  ROOT %add.0 = f32[256]{0} add(%param_0.7, %param_1.7), metadata={op_name="jit(program)/add" stack_frame_id=12}
}

ENTRY %main.4 (oh_arg.1: f32[256,32], idx_arg.1: s32[256]) -> f32[256] {
  %oh_arg.1 = f32[256,32]{1,0} parameter(0), metadata={op_name="oh_arg"}
  %idx_arg.1 = s32[256]{0} parameter(1), metadata={op_name="idx_arg"}
  %constant.2 = f32[] constant(0)
  %constant.3 = f32[32,16]{1,0} constant({...})
  %compare_convert_fusion = f32[256,32]{1,0} fusion(%idx_arg.1), kind=kLoop, calls=%fused_computation.2, metadata={op_name="jit(program)/onehot_onthefly/jit(_one_hot)/convert_element_type" stack_frame_id=9}
  %ynn_fusion.1 = f32[256]{0} fusion(%constant.2, %oh_arg.1, %constant.3), kind=kCustom, calls=%fused_computation.1, metadata={op_name="jit(program)/onehot_materialized/reduce_sum" stack_frame_id=4}, backend_config={"fusion_config":{"kind":"__ynn_fusion"},"outer_dimension_partitions":[]}
  %ynn_fusion = f32[256]{0} fusion(%constant.2, %compare_convert_fusion, %constant.3), kind=kCustom, calls=%fused_computation, metadata={op_name="jit(program)/onehot_onthefly/reduce_sum" stack_frame_id=6}, backend_config={"fusion_config":{"kind":"__ynn_fusion"},"outer_dimension_partitions":[]}
  ROOT %wrapped_add = f32[256]{0} fusion(%ynn_fusion.1, %ynn_fusion), kind=kLoop, calls=%wrapped_add_computation, metadata={op_name="jit(program)/add" stack_frame_id=12}
}

