WT talk 04.09.2024

ragged tensors and shape types

Question from last time I talked about ML and https://github.com/Mikolaj/horde-ad at WT:

does shape-typing (or rank-typing) of regular array ops limit how much we can express?

  op :: ArrayD Double -> ArrayR Int 3 -> ArrayS Char '[7, 4, 5]

  sumOuter :: ArrayD Int -> ArrayD Int
  sumOuter :: ArrayR Int (n + 1) -> ArrayR Int n
  sumOuter :: ArrayS Int (k ': sh) -> ArrayS Int sh

Answer: yes, but.

1. let's generalize from regular to jagged arrays (with offsets)
1.1. it's *not* sparse vs dense arrays (how many zeros, an implementation issue)

000200040002
002000400600
030000000000

1.2. jagged vs regular arrays (not just implementation but extra structure)

20x4x100

243255
   234
 2244444623463463
      22

1.3. jagged is complex, costly, less standard, less support
1.4. so there's merit in limiting oneself to regular arrays
1.5. but, alas, jaggedness emerges naturally from equalities/rewrites!

  build k (\i -> append a(i) b(i))
  == transpose (append (transpose (build k (\i -> a)))
                       (transpose (build k (\i -> b))))

  build 3 (\i -> append (replicate (i + 2) 'a') (replicate (5 - i) 'b')
  == transpose (append (transpose (build 3 (\i -> replicate (i + 2) 'a')))
                       (transpose (build 3 (\i -> replicate (5 - i) 'b'))))

aa
aaa
aaaa

bbbbb
bbbb
bbb

aaa
aaa
baa
bba
bbb
bbb
bbb

2. shape-typing of jagged array ops
2.1. shaped arrays, ArrayS Char '[7, 4, 5], immediately excludes ragged array ops,
     e.g., it bans elementwise (/= 0) filter
2.2. ranked arrays, ArrayR Int 3, depends on if "rank" means ragged or not and strength of the type system,
     e.g., the repeated dimensions filter op requires ranks dependent on values
2.3. "dynamic" arrays, ArrayD Double, anything goes
2.4. so the answer is "yes", but


3. of regular array ops
3.1. mismatch no longer possible between shape types and arrays structures
3.2. so it's all about type system strength vs the class of ops in question
3.3. the filter and dimension filter are still good examples to consider
3.4. so "it depends"


4. of regular array ops where result shape depends only on argument shapes
4.1. already for a modest simple-types-like system the answer is probably "no" (for recursive ops, etc.)
4.2. in horde-ad we ban dependent shapes not because of typing, but to avoid ragged arrays when rewriting as in 1.5
4.3. as a side-effect, our ranked and shaped arrays don't differ in expressiveness, 
     but only in how much they document and the amount of pain in proving shape vs rank equalities to GHC