fix: move shape definition before use in mean_dim function#21

Jah-yee wants to merge 1 commit into

thinking-machines-lab:mainthinking-machines-lab/batch_invariant_ops:mainfrom

Jah-yee:fix-shape-orderJah-yee/batch_invariant_ops:fix-shape-orderCopy head branch name to clipboard