Auto-match tile rank in load/store

[AntonOresten]

This is intended to reduce boilerplate reshaping by making trailing singletons implicit, i.e.:

q = ct.load(Q, (1, bid_x, head_idx, batch_idx), (D_K[], TILE_M[], 1, 1))
q = reshape(q, (D_K[], TILE_M[]))
...
acc = reshape(acc, (D_V[], TILE_M[], 1, 1))
ct.store(Out, (1, bid_x, head_idx, batch_idx), acc)

can become:

q = ct.load(Q, (1, bid_x, head_idx, batch_idx), (D_K[], TILE_M[]))
...
ct.store(Out, (1, bid_x, head_idx, batch_idx), acc)

There is precedence for this in that ct.store would already remove (and in specific cases add?) singletons. As per the current README:

Store reshaping

ct.store automatically reshapes the tile to match the target array's rank by dropping singleton dimensions (e.g., storing a (1, N) tile into a 1D array reshapes it to (N,)). Scalar () tiles are reshaped to (1,).

[maleadt]

I'm not sure I like this, at least not a generic shape-matching implementation. Dropping trailing singletons is probably fine, and we do this often in Julia. It's incompatible with the existing store shape matching which drops leading singletons, which is fishy. That said, for most of these ops (ct. prefixed one) I'd like to stay as close as possible to Python semantics.

[maleadt]

Pushed a commit that restricts the matching to consecutive trailing singletons, which matches Julia better. Thoughts?

[AntonOresten]

Yes, I think that's much better.

Did dropping leading singletons on store come from Python?

[maleadt]

Did dropping leading singletons on store come from Python?

No, I added it after making reductions preserve singleton dimensions (as Julia does), resulting in lots of mismatches between. But by allowing dropping singleton trailing dimensions we can still get rid of the dropdims previously needed by switching layouts around.