ordinal_tensor_1 = torch.tensor([index], dtype=torch.float).to(device)
  ordinal_tensor_2 = torch.tensor([index], dtype=torch.int32).to(device)

  xm.all_reduce(xm.REDUCE_SUM, [ordinal_tensor_1, ordinal_tensor_2])

ENTRY %IrToHlo.27 (p0.1: f32[], p1.2: s32[1], p2.3: f32[1]) -> (f32[1], s32[1]) {
.......
  %all-reduce.12 = (s32[1]{0}, s32[]) all-reduce(s32[1]{0} %get-tuple-element.6, s32[] %get-tuple-element.7), replica_groups={}, constrain_layout=true, to_apply=%AddComputation.8, metadata={op_type="xla__cross_replica_sum" op_name="xla__cross_replica_sum" source_file="/workspaces/dk3/pytorch/xla/torch_xla/core/xla_model.py" source_line=501}
.....
  %get-tuple-element.24 = f32[1]{0} get-tuple-element((f32[1]{0}, f32[]) %all-reduce.23), index=0, metadata={op_type="xla__cross_replica_sum" op_name="xla__cross_replica_sum" source_file="/workspaces/dk3/pytorch/xla/torch_xla/core/xla_model.py" source_line=501}
  %get-tuple-element.13 = s32[1]{0} get-tuple-element((s32[1]{0}, s32[]) %all-reduce.12), index=0, metadata={op_type="xla__cross_replica_sum" op_name="xla__cross_replica_sum" source_file="/workspaces/dk3/pytorch/xla/torch_xla/core/xla_model.py" source_line=501}
  ROOT %tuple.26 = (f32[1]{0}, s32[1]{0}) tuple(f32[1]{0} %get-tuple-element.24, s32[1]{0} %get-tuple-element.13)
}