scale parsed as float in ONNX scaled_dot_product_attention implementation

[nicolas-mng]

Hi,
I have been getting errors looking like the one below when trying to export a model to ONNX within which I manually provide a scale argument to the scaled dot product attention calls:

   File "/usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_opset14.py", line 176, in scaled_dot_product_attention
    query_scaled = g.op("Mul", query, g.op("Sqrt", scale))
   File "<@beartype(torch.onnx._internal.jit_utils.GraphContext.op) at 0x7f10f6b3feb0>", line 44, in op
 beartype.roar.BeartypeCallHintParamViolation: Method torch.onnx._internal.jit_utils.GraphContext.op() parameter raw_args=0.03125 violates type hint typing.Union[torch.Tensor, torch.Value], as float 0.03125 not <protocol "torch.Tensor"> or <class "torch.Value">.

Looking at symbolic_opset14.py, I see that _maybe_get_const parses scale as a float which triggers the error as the type is not correct, cf.
here

Is this intentional? If I remove this line, then I manage to export successfully my model.
Thanks

[thiagocrepaldi]

Propose a PR and provide a unit test and we can review it for you