coremltools
coremltools copied to clipboard
Patch convert nodes op to allow for ops with private variable names
This PR addresses an issue with the convert_nodes method that arises for ops with privately named variables such as __and__
Convert_nodes assumes that ops ending in _ are in place ops and removes the trailing _ before op lookup.
op_lookup = node.kind
if op_lookup.endswith("_"):
# This is an "in place" op.
# Look up the standard op instead by removing underscore.
op_lookup = op_lookup[:-1]
add_op = _TORCH_OPS_REGISTRY.get(op_lookup, None)
the patch below fixes this
add_op = _TORCH_OPS_REGISTRY.get(op_lookup, None)
if add_op is None and op_lookup.endswith("_"):
# This is an "in place" op.
# Look up the standard op instead by removing underscore.
op_lookup = op_lookup[:-1]
add_op = _TORCH_OPS_REGISTRY.get(op_lookup, None)
I mentioned this previously here https://github.com/apple/coremltools/pull/1509#discussion_r892186687
@dncnbuck - could you please add a unit test which fails without your fix but passes with the fix?