Support stash_type attribute for onnx.LayerNormalization

[jinchen62]

Fixes nod-ai/SHARK-ModelDev#888

If stash_type is different from input_dtype/result_dtype:

1. convert x dtype to stash_type
2. calculate mean and var in stash_type since x is in stash_type already
3. convert back to result_dtype before stage two calculation
4. convert mean_dtype and var_dtype if they are different from stash_type

[zjgarvey]

I think we should probably support the stash type arg by separating the two stages of computation as is suggested by ONNX in https://onnx.ai/onnx/operators/onnx__LayerNormalization.html. If an onnx op actually has different result types and stash types, we would likely see numeric mismatches for those situations unless we perform the computation correctly.

Another option is to allow LayerNormalization to be function-expanded on import via

torch-mlir/python/torch_mlir/extras/onnx_importer.py

Line 102 in 1b8d7e0

function_expansion_allowlists_by_domain: Optional[Dict[str, set[str]]] = field(

In any case, we should put together a few e2e tests for this op:

1. With bf16 result type and bf16 stash type
2. With bf16 result type and unspecified stash type