Проблема несоответствия размеров Futhark

Question

Я пытаюсь реализовать групповую функциональность. Я пришел к выводу, что мой код должен быть правильным.

Вот важный раздел

let type_func (typ: i32) (v1 : u32)  (v2: u32) : u32 =
        match typ
        case 1 -> (*) v1 v2
        case 2 -> (+) v1 v2
        case 3 -> (u32.max) v1 v2
        case 4 -> (u32.min) v1 v2
        case x -> (u32.min) v1 v2 -- TODO change to some panic function


let merge [n][m] (s_cols_t: [n]i32) (a: [m]u32) (b: [m]u32) : [m]u32 =
  map (\i -> if i == 0 then a[i]
                       else (type_func s_cols_t[i-1] a[i] b[i])
      ) (iota m)

let main [n][m][t] (db : [n][m]u32)  (g_col: i32) (s_cols: [t]i32) (t_cols: [t]i32) : [][]u32 =
  let keep_g = db[:, g_col]
  let keep_s_cols = map (\c -> db[:, c]) s_cols
  let keep_inter = concat [keep_g] (keep_s_cols)
  let keep = transpose keep_inter
  let sorted_rows = rsort keep -- ideally pass groupby col here
  let idxs = mk_flags sorted_rows[:, 0]
  let flag = map (== 1) idxs
  let helper = merge t_cols
  in segmented_reduce helper (replicate (length keep_inter) 0)  flag sorted_rows

Однако компилятор выдает следующую ошибку.

[0]> :l groupby.fut
Loading groupby.fut
Error at groupby.fut:63:70-80 :
Cannot apply "segmented_reduce" to "sorted_rows" (invalid type).
Expected: [n][argdim₃₅]u32
Actual:   *[n][ret₁₃]u32

Dimensions "argdim₃₅" and "ret₁₃" do not match.

Note: "argdim₃₅" is value of argument
        length keep_inter
      passed to "replicate" at 63:42-58.

Note: "ret₁₃" is unknown size returned by "concat" at 57:20-48.

У меня есть вручную проверил, что измерение argdim₃₅ и ret₁₃ совпадает, вставляя код в строку в REPL. Это просто ограничение компилятора, с которым я столкнулся, или я делаю что-то глупое?

Answer 1

У меня есть хакерское решение проблемы. Вы можете поместить большую часть основных функций в другую функцию, а затем явно задать форму измерения, которое используется в определении функции.

let groupby [n][m][s][t] (db : [n][m]u32)  (cols: [s]i32)  (t_cols: [t]i32) : [][]u32 =
  let keep_fun columns row = map (\i -> row[i]) columns
  let keep = map (keep_fun cols) db
  let sorted_rows = rsort keep -- ideally pass groupby col here
  let idxs = mk_flags sorted_rows[:, 0]
  let flag = map (== 1) idxs
  let helper = merge t_cols
  in segmented_reduce helper (replicate s 0)  flag sorted_rows

let main db g_col s_cols t_cols =
  let cols = concat [g_col] s_cols
  in groupby db cols t_cols

Это чрезвычайно варварское c и уродливое решение, но если это работает, это работает.