Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Account for additional cases in Linear node conversion #709

Merged
merged 2 commits into from
Aug 28, 2023

Conversation

antimora
Copy link
Collaborator

Pull Request Template

Checklist

  • Confirm that run-checks script has been executed.

Changes

Account for two additional cases in Linear node conversion. Previously we only converted Gemm node to a Linear module. Now, we are converting MatMul and Add optionally. See ONNX graph of 3 cases for illustration.

Testing

Add more test cases and ONNX example.

image

Copy link
Member

@nathanielsimard nathanielsimard left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A small refactor for cleaner code, but otherwise great !

Comment on lines 148 to 157
if peek_node.is_some()
&& peek_node.unwrap().node_type == NodeType::Add
&& peek_node.unwrap().inputs.len() == 2

// Make sure the Add node has a value in one of its inputs and
// the other input is the output of this MatMul node
&& (peek_node.unwrap().inputs[0].name == node.outputs[0].name
&& peek_node.unwrap().inputs[1].value.is_some())
| (peek_node.unwrap().inputs[1].name == node.outputs[0].name
&& peek_node.unwrap().inputs[0].value.is_some())
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Convert the condition into a function. You can early return if peak_node.is_none(), then call unwrap only once.

@nathanielsimard nathanielsimard merged commit 13342b4 into tracel-ai:main Aug 28, 2023
@nathanielsimard nathanielsimard deleted the onnx-linear-fix branch August 28, 2023 20:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants