[jit] Support in membership checks
#21527
Conversation
This PR adds support for `in` checks like `2 in [1, 2, 3]` or `key in my_dict` For objects it uses the magic method `__contains__(self, key)`
pytorchbot
added
oncall: jit
torch/csrc/jit/script/parser.cpp
Outdated
| if (kind == TK_FOR) { | ||
| in_for = true; | ||
| auto target = parseExp(); | ||
| in_for = false; |
There was a problem hiding this comment.
This doesn't work for something like: for i in range(d["a"] if "a" in d else 2):
zdevito
left a comment
zdevito
left a comment
There was a problem hiding this comment.
I think there is one bug in the implementation and the parser needs to be updated to distinguish variable declarations in for loops from expressions.
pytorchbot
added
the
module: pybind
zdevito
left a comment
zdevito
left a comment
There was a problem hiding this comment.
Looks good. Just one minor suggestion.
| auto named_values = getNamedValues(inputs, /*maybe_unpack=*/false); | ||
| auto class_arg_position = 0; | ||
|
|
||
| if (tree->kind() == TK_IN) { |
There was a problem hiding this comment.
The class_arg_position is weird here. This would be the appropriate time to flip the order and call the method. Previously it was invalid because the subexpressions were not yet emitted. But now they are (getNamedValues does it) and it would fit better in the API if it just called std::swap on named_values.
|
|
||
| if (kind == TK_FOR) { | ||
| auto target = parseExp(); | ||
| auto target = parseIdent(); |
There was a problem hiding this comment.
why targets will be a Ident instead of Expr? things like a, (b, c) = 1, (2, 3) will not work. Why do you need to change it in the first place?
There was a problem hiding this comment.
We don't support that currently so this isn't a regression at least (we just expect everything to be an Ident somewhere later in the compilation process. Expr includes in statements now to keep it from parsing the in on a for-in loop we need to parse the loop exprlist as just identifiers.
Summary: This PR adds support for `in` checks like `key in my_dict` For now it leaves lists as a follow up due to the changes around `IValue` lists and it needing an `IValue` equality op. For objects it uses the magic method `__contains__(self, key)` ](https://our.intern.facebook.com/intern/diff/15811203/) Pull Request resolved: pytorch/pytorch#21527 Pulled By: driazati Differential Revision: D15811203 fbshipit-source-id: 95745060394f8a9450efaaf8ab09d9af83bea01e
This PR adds support for
inchecks likekey in my_dictFor now it leaves lists as a follow up due to the changes around
IValuelists and it needing anIValueequality op.For objects it uses the magic method
__contains__(self, key)Differential Revision: D15811203