coremltools icon indicating copy to clipboard operation
coremltools copied to clipboard

Patch convert nodes op to allow for ops with private variable names

Open dncnbuck opened this issue 3 years ago • 1 comments

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 avatar Sep 03 '22 09:09 dncnbuck

@dncnbuck - could you please add a unit test which fails without your fix but passes with the fix?

TobyRoseman avatar Sep 07 '22 17:09 TobyRoseman