embedding_to_f16.mlir

module {
  func.func @decode_bs4(%arg0: !torch.vtensor<[4,?],si64>, %arg1: !torch.vtensor<[128256,4096],f16>) -> !torch.vtensor<[4,?,4096],f32> {
    %false = torch.constant.bool false
    %false_0 = torch.constant.bool false
    %int-1 = torch.constant.int -1
    %int6 = torch.constant.int 6
    %0 = torch.prims.convert_element_type %arg1, %int6 : !torch.vtensor<[128256,4096],f16>, !torch.int -> !torch.vtensor<[128256,4096],f32>
    %1 = torch.aten.embedding %0, %arg0, %int-1, %false_0, %false : !torch.vtensor<[128256,4096],f32>, !torch.vtensor<[4,?],si64>, !torch.int, !torch.bool, !torch.bool -> !torch.vtensor<[4,?,4096],f32>
    return %1 : !torch.vtensor<[4,?,4096],f32>
  }
}