Handling bfloat16 constant propagation

[gramalingam]

Just a temporary workaround for #1471 for experimentation/discussion.

It looks like the issue is that bfloat16 tensor constants are represented as float32 numpy arrays (in ONNX itself), when converted to numpy array. In the context of constant-propagation, this means that we cannot rely solely on the numpy value's dtype to figure out the ONNX type.

The hack below suppresses constant-propagation for bfloat16 constants: partially because of the above reason, and partially since I am yet unclear if this convention is supported by the onnx reference implementation (or ORT), etc. Assuming the backend supports it, we can try other alternative solutions too. One possibility is to simply suppress constant-propagation if the output-types are unknown (in the onnx model).

[codecov]

Codecov Report

Attention: Patch coverage is 33.33333% with 8 lines in your changes are missing coverage. Please review.

Project coverage is 77.59%. Comparing base (ac7ce49) to head (cad2e33).

Files	Patch %	Lines
onnxscript/optimizer/constant_folding.py	33.33%	3 Missing and 3 partials ⚠️
onnxscript/_legacy_ir/visitor.py	33.33%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1484      +/-   ##
==========================================
- Coverage   77.63%   77.59%   -0.05%     
==========================================
  Files         212      212              
  Lines       23011    23023      +12     
  Branches     3945     3951       +6     
==========================================
- Hits        17865    17864       -1     
- Misses       4393     4399       +6     
- Partials      753      760       +7     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[github-actions]

Test Results

     28 files  ±     0      28 suites  ±0   2h 37m 28s ⏱️ - 14m 3s
  5 897 tests + 2 024   4 005 ✅ +1 614    1 880 💤 +   405  11 ❌ + 4  1 🔥 +1 
492 370 runs   - 34 010  79 359 ✅  - 2 735  412 921 💤  - 31 252  89 ❌  - 24  1 🔥 +1 

For more details on these failures and errors, see this check.

Results for commit 7152688. ± Comparison against base commit 7b10687.

This pull request removes 4 and adds 2028 tests. Note that renamed tests count towards both.

tests.ir.serde_test.SerdeTest ‑ test_serialization_deserialization_produces_same_model_01_torchscript_model_onnx
tests.ir.serde_test.SerdeTest ‑ test_serialization_deserialization_produces_same_model_02_mobilenetv2_100_dynamo_onnx
tests.ir.serde_test.SerdeTest ‑ test_serialization_deserialization_produces_same_model_03_resnet18_dynamo_onnx
tests.ir.serde_test.SerdeTest ‑ test_serialization_deserialization_produces_same_model_04_Speech2Text2ForCausalLM_dynamo_onnx

onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_basic_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_doc_string
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_if_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_if_loop_defs
onnxscript._internal.analysis_test.TestAssignedVarAnalysis ‑ test_loop_defs
onnxscript._internal.analysis_test.TestExposedUses ‑ test_basic
onnxscript._internal.analysis_test.TestExposedUses ‑ test_called_function
onnxscript._internal.analysis_test.TestExposedUses ‑ test_doc_string
onnxscript._internal.analysis_test.TestExposedUses ‑ test_for_loop
onnxscript._internal.analysis_test.TestExposedUses ‑ test_if
…