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

Export control-flow Ops in ppe.onnx.export #648

Merged
merged 17 commits into from
Mar 8, 2023

Conversation

HiroakiMikami
Copy link
Member

@HiroakiMikami HiroakiMikami commented Mar 5, 2023

Add functions to represent control flows in jax.lax (fori_loop, while_loop, cond)

These functions trace only first iteration, then ppe.onnx inserts onnx::Loop or onnx::If operators to an exported ONNX as postprocess.

@HiroakiMikami
Copy link
Member Author

/test

@HiroakiMikami HiroakiMikami changed the title [WIP] control-flow Export control-flow Ops in ppe.onnx.export Mar 6, 2023
@take-cheeze
Copy link
Contributor

/test


@contextmanager
def init_lax_state() -> Generator[None, None, None]:
_lax_state.n_call = 0
Copy link
Contributor

Choose a reason for hiding this comment

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

memo: This attribute is to avoid duplicate name created by as_output

Copy link
Contributor

@take-cheeze take-cheeze left a comment

Choose a reason for hiding this comment

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

Basic functionallity was LGTM 👍
Left few suggestion comment what I fealt reading code

@HiroakiMikami
Copy link
Member Author

/test

@HiroakiMikami
Copy link
Member Author

I reflected your comments 🙇 @take-cheeze

Copy link
Contributor

@take-cheeze take-cheeze left a comment

Choose a reason for hiding this comment

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

Thank you! 👍

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants