Add additional argument to _aten_bmm

[jrplatin]

Fixes this error from running MLA + TorchAX:

(EngineCore_DP0 pid=613611)   File "/mnt/disks/jacobplatin/vllm/vllm/model_executor/layers/attention/mla_attention.py", line 440, in forward
(EngineCore_DP0 pid=613611)     self.forward_impl(
(EngineCore_DP0 pid=613611)   File "/mnt/disks/jacobplatin/vllm/vllm/model_executor/layers/attention/mla_attention.py", line 614, in forward_impl
(EngineCore_DP0 pid=613611)     torch.bmm(mqa_q_nope, self.W_UK_T, out=mqa_ql_nope)
(EngineCore_DP0 pid=613611)   File "/mnt/disks/jacobplatin/torchax/torchax/tensor.py", line 254, in __torch_function__
(EngineCore_DP0 pid=613611)     return func(*args, **(kwargs or {}))
(EngineCore_DP0 pid=613611)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=613611)   File "/mnt/disks/jacobplatin/torchax/torchax/tensor.py", line 277, in __torch_dispatch__
(EngineCore_DP0 pid=613611)     return self.env.dispatch(func, types, args, kwargs)
(EngineCore_DP0 pid=613611)            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=613611)   File "/mnt/disks/jacobplatin/torchax/torchax/tensor.py", line 599, in dispatch
(EngineCore_DP0 pid=613611)     res = op.func(*args, **kwargs)
(EngineCore_DP0 pid=613611)           ^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore_DP0 pid=613611) TypeError: _aten_bmm() got an unexpected keyword argument 'out'

[wdhongtw]

Maybe add an assertion to ensure the out is indeed None?

[wdhongtw]

Also, there is a helper class OutVariant (and a mapping _out_variant_to_functional), not sure if it can solve your problem.