Batched multi_dot / chain_matmul + let it accept a tensor instead of tuple #55261
Description
This is useful for computing an accumulated transition matrix from individual transition matrices:
- First usecase: (BxN1xN2, BxN2xN3, ...) -> BxN1xN3
- Second usecase: (BxTxNxN, dim = 1) -> BxNxN -> this could save an unbind call and maybe a GPU sync (could also be (TxBxNxN, dim = 0) -> BxNxN)
Also may be useful to allow this for logmm in addition to mm:
In the wild this happens in https://github.com/ajabri/videowalk/blob/master/code/model.py#L149 (computing the cycle's total transition matrix)
Originally requested here:
cc @jianyuh @nikitaved @pearu @mruberry @heitorschueroff @walterddr @IvanYashchuk