Use ONNX Rewriter and IR to simplify the mnb_to_qdq pass #1482
Conversation
Co-authored-by: Jambay Kinley <[email protected]>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
|
||
| # Add Logic handling input 3 | ||
|
|
||
| unpacked_weight_arrays = _unpack_weights( |
Check failure
Code scanning / lintrunner
RUFF/F821 Error
| @@ -7,8 +7,10 @@ | |||
| from pathlib import Path | |||
| from typing import TYPE_CHECKING, Any, Dict | |||
|
|
|||
| import ml_dtypes | |||
Check notice
Code scanning / CodeQL
Unused import Note
| @@ -7,8 +7,10 @@ | |||
| from pathlib import Path | |||
| from typing import TYPE_CHECKING, Any, Dict | |||
|
|
|||
| import ml_dtypes | |||
Check warning
Code scanning / lintrunner
PYLINT/W0611 Warning
| @@ -7,8 +7,10 @@ | |||
| from pathlib import Path | |||
| from typing import TYPE_CHECKING, Any, Dict | |||
|
|
|||
| import ml_dtypes | |||
Check warning
Code scanning / lintrunner
RUFF/F401 Warning
| matmul = op.Add(matmul, bias) | ||
| return matmul | ||
|
|
||
| replace_mat_mul_n_bits = orp.RewriteRule( |
Check notice
Code scanning / CodeQL
Unused local variable Note
| return False | ||
| g_idx = g_idx.constant_value.numpy() | ||
| trivial_g_idx = np.arange(k, dtype=np.int32) // block_size | ||
| if not np.array_equal(g_idx, trivial_g_idx): |
Check warning
Code scanning / lintrunner
RUFF/SIM103 Warning
| g_idx = g_idx.constant_value.numpy() | ||
| trivial_g_idx = np.arange(k, dtype=np.int32) // block_size | ||
| if not np.array_equal(g_idx, trivial_g_idx): | ||
| # TODO: We can log why the pattern is not matched here |
Check warning
Code scanning / lintrunner
RUFF/TD002 Warning
| matmul = op.Add(matmul, bias) | ||
| return matmul | ||
|
|
||
| replace_mat_mul_n_bits = orp.RewriteRule( |
Check warning
Code scanning / lintrunner
PYLINT/W0612 Warning
| matmul = op.Add(matmul, bias) | ||
| return matmul | ||
|
|
||
| replace_mat_mul_n_bits = orp.RewriteRule( |
Check warning
Code scanning / lintrunner
RUFF/F841 Warning
| graph: ir.Graph = context.graph | ||
| return value in graph.initializers.values() | ||
|
|
||
| def mat_mul_n_bits_pattern_check(context, *, q_weight, g_idx, mat_mul_n_bits_out: ir.Value, **_) -> bool: |
There was a problem hiding this comment.
does q_weight here match for the input right before g_idx or it is whatever it is in the mat_mul_n_bits_pattern signature? The input before g_idx is qzero and can be optional. we want to check the second input
There was a problem hiding this comment.
The inputs of the pattern-function (mat_mul_n_bits_pattern) are bound to values in the graph, and these values are passed in as keyword-arguments to the rewrite function here. So, the order here doesn't really matter, though I usually just copy-paste and use the same argument list for both.
| del node.meta["N"] | ||
|
|
||
| # TODO(justinchuby): Register and remove initializers | ||
| ir_model.opset_imports[""] = max(21, ir_model.opset_imports[""]) |
There was a problem hiding this comment.
TODO: Use a more robust version conversion process
Describe your changes
Checklist before requesting a review
lintrunner -a(Optional) Issue link