Skip to main content

Clashluke's group workspace

Timestamps visible
2022-10-30 14:17:45
     Unpadded size: 64.00M
2022-10-30 14:17:45
     XLA label: broadcast.7198.clone.2 = broadcast(get-tuple-element.18418), dimensions={}
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  9. Size: 32.00M
2022-10-30 14:17:45
     Shape: bf16[4096,1,4096]{2,0,1:T(8,128)(2,1)}
2022-10-30 14:17:45
     Unpadded size: 32.00M
2022-10-30 14:17:45
     XLA label: copy.785 = copy(get-tuple-element.18154)
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  10. Size: 32.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/broadcast_in_dim[shape=(16, 1, 1, 256) broadcast_dimensions=()]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/model/main.py" source_line=60
2022-10-30 14:17:45
     Shape: f32[16,1,1,256]{3,2,1,0:T(2,128)}
2022-10-30 14:17:45
     Unpadded size: 16.0K
2022-10-30 14:17:45
     Extra memory due to padding: 16.0K (2.0x expansion)
2022-10-30 14:17:45
     XLA label: broadcast.7199 = broadcast(get-tuple-element.18418), dimensions={}
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  11. Size: 32.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/broadcast_in_dim[shape=(16, 1, 1, 256) broadcast_dimensions=()]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/model/main.py" source_line=60
2022-10-30 14:17:45
     Shape: f32[16,1,1,256]{3,2,1,0:T(2,128)}
2022-10-30 14:17:45
     Unpadded size: 16.0K
2022-10-30 14:17:45
     Extra memory due to padding: 16.0K (2.0x expansion)
2022-10-30 14:17:45
     XLA label: broadcast.7199.clone = broadcast(get-tuple-element.18418), dimensions={}
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  12. Size: 32.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/broadcast_in_dim[shape=(16, 1, 1, 256) broadcast_dimensions=()]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/model/main.py" source_line=60
2022-10-30 14:17:45
     Shape: f32[16,1,1,256]{3,2,1,0:T(2,128)}
2022-10-30 14:17:45
     Unpadded size: 16.0K
2022-10-30 14:17:45
     Extra memory due to padding: 16.0K (2.0x expansion)
2022-10-30 14:17:45
     XLA label: broadcast.7199.clone.1 = broadcast(get-tuple-element.18418), dimensions={}
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  13. Size: 32.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/broadcast_in_dim[shape=(16, 1, 1, 256) broadcast_dimensions=()]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/model/main.py" source_line=60
2022-10-30 14:17:45
     Shape: f32[16,1,1,256]{3,2,1,0:T(2,128)}
2022-10-30 14:17:45
     Unpadded size: 16.0K
2022-10-30 14:17:45
     Extra memory due to padding: 16.0K (2.0x expansion)
2022-10-30 14:17:45
     XLA label: broadcast.7199.clone.2 = broadcast(get-tuple-element.18418), dimensions={}
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  14. Size: 4.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/cond/branch_1_fun/while/body/while/body/cond[linear=(False, False, False)]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/shampoo.py" source_line=148
2022-10-30 14:17:45
     Shape: (f32[256,256]{1,0:T(8,128)}, f32[256,256]{1,0:T(8,128)})
2022-10-30 14:17:45
     Unpadded size: 4.0K
2022-10-30 14:17:45
     XLA label: tuple.742 = tuple(copy.912, get-tuple-element.17140)
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  15. Size: 4.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/reduce_sum[axes=(0, 1)]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/optimizer.py" source_line=101
2022-10-30 14:17:45
     Shape: (f32[]{:T(256)}, f32[256,256]{1,0:T(8,128)})
2022-10-30 14:17:45
     Unpadded size: 4.0K
2022-10-30 14:17:45
     XLA label: fusion.1869 = fusion(get-tuple-element.18143, add.2502, multiply.7882, get-tuple-element.18300, ...(+2)), kind=kLoop, calls=fused_computation.858.clone
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  16. Size: 4.0K
2022-10-30 14:17:45
     Shape: u32[8,128]{1,0}
2022-10-30 14:17:45
     Unpadded size: 4.0K
2022-10-30 14:17:45
     XLA label: constant literal
2022-10-30 14:17:45
     Allocation type: global
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  17. Size: 4.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/jvp(while)/body/cond/branch_0_fun/dot_general[dimension_numbers=(((2,), (0,)), ((), ())) precision=(<Precision.HIGHEST: 2>, <Precision.HIGHEST: 2>) preferred_element_type=None]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/model/reversible.py" source_line=78
2022-10-30 14:17:45
     Shape: (f32[2048,512]{1,0:T(8,128)}, bf16[2048,256,512]{2,1,0:T(8,128)(2,1)})
2022-10-30 14:17:45
     Unpadded size: 4.0K
2022-10-30 14:17:45
     XLA label: fusion.45 = fusion(get-tuple-element.386, get-tuple-element.13796, reduce.955, fusion.48), kind=kOutput, calls=fused_computation.45
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  18. Size: 4.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/jvp(while)/body/cond/branch_0_fun/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=(<Precision.HIGHEST: 2>, <Precision.HIGHEST: 2>) preferred_element_type=None]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/model/reversible.py" source_line=78
2022-10-30 14:17:45
     Shape: (f32[2048,512]{1,0:T(8,128)}, bf16[2048,256,512]{2,1,0:T(8,128)(2,1)})
2022-10-30 14:17:45
     Unpadded size: 4.0K
2022-10-30 14:17:45
     XLA label: fusion.40 = fusion(bitcast, get-tuple-element.384), kind=kOutput, calls=fused_computation.40
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  19. Size: 4.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/jvp(while)/body/cond/branch_1_fun/dot_general[dimension_numbers=(((1,), (0,)), ((), ())) precision=(<Precision.HIGHEST: 2>, <Precision.HIGHEST: 2>) preferred_element_type=None]" source_file="/home/ubuntu/HomebrewNLP-Jax/src/model/reversible.py" source_line=78
2022-10-30 14:17:45
     Shape: (f32[2048,512]{1,0:T(8,128)}, bf16[2048,256,512]{2,1,0:T(8,128)(2,1)})
2022-10-30 14:17:45
     Unpadded size: 4.0K
2022-10-30 14:17:45
     XLA label: fusion.57 = fusion(bitcast.6, get-tuple-element.442), kind=kOutput, calls=fused_computation.56
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================
2022-10-30 14:17:45
  20. Size: 4.0K
2022-10-30 14:17:45
     Operator: op_name="pmap(jitless_step)/jit(main)/while/body/jvp(while)/body/dynamic_update_slice" source_file="/home/ubuntu/HomebrewNLP-Jax/src/model/main.py" source_line=60
2022-10-30 14:17:45
     Shape: (bf16[16,256,4096,256]{1,3,2,0:T(8,128)(2,1)}, bf16[256,4096,256]{0,2,1:T(8,128)(2,1)})
2022-10-30 14:17:45
     Unpadded size: 4.0K
2022-10-30 14:17:45
     XLA label: fusion.1525 = fusion(get-tuple-element.17243, select.867, fusion.1485, fusion.1506, ...(+4)), kind=kLoop, calls=fused_computation.267.clone
2022-10-30 14:17:45
     Allocation type: HLO temp
2022-10-30 14:17:45
     ==========================