Accept inputs shape in Attention/MLA constructor#2089

copybara-service[bot] merged 1 commit intomainAI-Hypercomputer/maxtext:mainfrom

bvandermoon-nnx-attention-shapeAI-Hypercomputer/maxtext:bvandermoon-nnx-attention-shapeCopy head branch name to clipboard