ConstantFolder fails shape inference on second iteration

[thiagocrepaldi]

During ONNX export process, onnxscript.optimizer.optimize is called, which runs 2 times applying all optimizations.
When the optimization modified = fold_constants(model, external_data_folder, onnx_shape_inference=onnx_shape_inference) is executed TWICE, the following shape inference error is raised

Traceback (most recent call last):
  File "/opt/pytorch/torch/onnx/_internal/exporter.py", line 1509, in dynamo_export
    ).export()
      ^^^^^^^^
  File "/opt/pytorch/torch/onnx/_internal/exporter.py", line 1280, in export
    onnx_model = optimizer.optimize(onnx_model)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/onnxscript/onnxscript/optimizer/__init__.py", line 61, in optimize
    model = onnx.shape_inference.infer_shapes(
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/ptca/lib/python3.11/site-packages/onnx/shape_inference.py", line 46, in infer_shapes
    inferred_model_str = C.infer_shapes(
                         ^^^^^^^^^^^^^^^
onnx.onnx_cpp2py_export.shape_inference.InferenceError: [ShapeInferenceError] (op_type:MatMul, node name: torch_nn_modules_linear_Linear_fc1_1_0_aten_mm_1_n0): B has inconsistent type tensor(float)

Repro option 1:

class MyLinear(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 5, bias=False, dtype=torch.bfloat16)

    def forward(self, tensor_x: torch.Tensor):
        return self.fc1(tensor_x)

model = MyLinear()
tensor_x = torch.rand((3, 10), dtype=torch.bfloat16)
onnx_program = torch.onnx.dynamo_export(model, tensor_x)

Repro option 2:

use the onnx model attached to the issue and run the pass manually:

onnx_model.zip

import onnx
import onnxscript
model = onnx.load("/path/to/model.onnx")
_ = onnxscript.optimizer.fold_constants(model,"/path/to
/", onnx_shape_inference=True)
model = onnx.shape_inference.infer_shapes(model, check_type=True, strict_mode=True, data_prop=True)
_ = onnxscript.optimizer.fold_constants(model,"/opt/dev/no_op/", onnx_shape_inference=True)
model = onnx.shape_inference.infer_shapes(model, check_type=True, strict_mode=True, data_prop=True)

[gramalingam]

I stepped through this. It looks like the constant-propagation evaluates the Transpose of the initializer. The loaded initializer is a float32 numpy array (even though it is declared to be a bfloat16 initializer). @thiagocrepaldi , do you know if the initializer is stored externally as a float32 tensor?

[gramalingam]

However, it is possible that onnx helpers create a numpy float32 array (for bfloat16). This could be one source of the issue.

[gramalingam]

The fix in PR #1484 is a partial one. For a more complete fix, we need to consider the following:

• Does the backend implementation used for constant-propagation (currently, ONNX's reference implementation, but could be onnxruntime also in future) support bfloat16 inputs/outputs? How are they represented? Eg., a float32 numpy array is insufficient to distinguish between float32 and a float32 representation of bfloat16. (Talking about backend conventions.)
• We could also adapt the fix in Handling bfloat16 constant propagation #1484 to do it on the value returned by the backend ... either based on the convention above, or using pre-existing types. That is, we could disable constant-prop unless we know the type of the return value. In other words, existing fix doesn't handle the case of ops where the input is not bfloat16, but output is.