ra2a_r8_2x4.hlo

ra2a_r8_2x4.hlo

HloModule module, input_output_alias={ {}: (1, {}) }, num_partitions=1, replica_count=8
// HloModule module, num_partitions=1, replica_count=8

ENTRY entry {
  id = u32[] replica-id()
  input = f32[16] parameter(0)
  output = f32[16] parameter(1)
  send_sizes = s32[8] constant({1,1,1,1,1,1,1,1})
  recv_sizes = s32[8] constant({1,1,1,1,1,1,1,1})
  input_offsets = s32[8] constant({0, 2, 4, 6, 8, 10, 12, 14})
  step = u32[] constant(2)
  oof = u32[] multiply(id, step)
  oof2 = s32[] convert(oof)
  output_offsets = s32[8] broadcast(oof2)

  ra2a = f32[16] ragged-all-to-all(input, output, input_offsets,
  send_sizes, output_offsets, recv_sizes), replica_groups={{0,1,2,3,4,5,6,7}}
  ROOT res = cosine(ra2a)
}