[JIT][script] Add matmul(@), pow(**) operator#7648
Conversation
| with self.assertRaisesRegex(NotSupportedError, "unsupported binary operator:"): | ||
| @torch.jit.script | ||
| def binop(x, y): | ||
| # Replace this with another unsupported op when/if it gets supported |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
apaszke
left a comment
There was a problem hiding this comment.
Mostly LGTM, but needs some adjustments!
| @unittest.skipIf(not PY35, "Python 3.5 needed") | ||
| def test_matmul(self): | ||
| def func(a, b): | ||
| return a @ b |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| with self.assertRaisesRegex(NotSupportedError, "unsupported binary operator:"): | ||
| @torch.jit.script | ||
| def binop(x, y): | ||
| # Replace this with another unsupported op when/if it gets supported |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
torch/csrc/jit/script/lexer.h
Outdated
| {'<', '>', TK_EQ, TK_LE, TK_GE, TK_NE}, | ||
| {'+', '-'}, | ||
| {'*', '/'}, | ||
| {'*', '/', TK_MATMUL, TK_POW}, |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
zdevito
left a comment
There was a problem hiding this comment.
Thanks! I have a few small bugs to fix, otherwise it is good.
torch/csrc/jit/script/lexer.h
Outdated
| _(TK_UNARY_MINUS, "unary minus", "") | ||
| _(TK_UNARY_MINUS, "unary minus", "") \ | ||
| _(TK_POW, "pow operator", "**") \ | ||
| _(TK_MATMUL, "matmul operator", "@") |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
torch/csrc/jit/script/lexer.h
Outdated
| {'<', '>', TK_EQ, TK_LE, TK_GE, TK_NE}, | ||
| {'+', '-'}, | ||
| {'*', '/'}, | ||
| {'*', '/', TK_MATMUL, TK_POW}, |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
torch/jit/frontend.py
Outdated
| ast.Mult: '*', | ||
| ast.Div: '/', | ||
| } | ||
| if PY2: |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
|
|
||
| def test_pow(self): | ||
| def func(a, b): | ||
| return a ** b |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
torch/csrc/jit/script/lexer.h
Outdated
| _(TK_UNARY_MINUS, "unary minus", "") | ||
| _(TK_UNARY_MINUS, "unary minus", "") \ | ||
| _(TK_POW, "pow operator", "**") \ | ||
| _(TK_MATMUL, "matmul operator", "@") |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
jamesr66a
left a comment
There was a problem hiding this comment.
Small comment about code duplication
test/test_jit.py
Outdated
| spec = importlib.util.spec_from_file_location('test_matmul_py3', script_path) | ||
| module = importlib.util.module_from_spec(spec) | ||
| spec.loader.exec_module(module) | ||
| fn = module.fn |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
* add matmul(@), pow(**) operator * fix bug(matmul not in py2) in @ operator * fix bugs * add get_fn help func to remove duplication in test_jit
Add matmul(@) operator (#6049)
Add pow(**) operator (#7542)