Unsupported tensorflow ops: `PartitionedCall`

[josephrocca]

Describe the bug
The BitwiseAnd and PartitionedCall ops are not supported.

Urgency
It's a blocker for me, but there are no specific deadline dates.

System information

• OS Platform and Distribution: Linux-5.4.144+-x86_64-with-Ubuntu-18.04-bionic (Google Colab)
• Tensorflow Version: 2.8.0
• Python version: 3.7.12

To Reproduce
Simply run this Colab notebook: https://colab.research.google.com/drive/1bwa0BbTgnV214EnWV0nyNGYspynZYf06

Additional context
Here's the command that I used:

!python -m tf2onnx.convert --saved-model /content/pixel_text_score --output pixel_text_score.onnx --opset 15 --verbose --continue_on_error

And here's the full logs:

/usr/lib/python3.7/runpy.py:125: RuntimeWarning: 'tf2onnx.convert' found in sys.modules after import of package 'tf2onnx', but prior to execution of 'tf2onnx.convert'; this may result in unpredictable behaviour
  warn(RuntimeWarning(msg))
2022-02-23 03:08:47,129 - WARNING - tf2onnx.tf_loader: '--tag' not specified for saved_model. Using --tag serve
tcmalloc: large alloc 1212309504 bytes == 0x55804ad54000 @  0x7f50396601e7 0x7f4fe2e6b90b 0x7f4fdcde9bf8 0x7f4fdcdddde2 0x558043e9ef78 0x558043f12a6d 0x558043f0d02f 0x558043e9faba 0x558043f0e108 0x558043e9f9da 0x558043f0deae 0x558043e9f9da 0x558043f122c0 0x558043f0d66e 0x558043e9faba 0x558043f0deae 0x558043f0d02f 0x558043e9faba 0x558043f0ecd4 0x558043f0d66e 0x558043e9faba 0x558043f0deae 0x558043f0d02f 0x558043e9faba 0x558043f0ecd4 0x558043e9f9da 0x558043f0deae 0x558043f0d02f 0x558043f0cd43 0x558043f0b1b0 0x558043e9e229
tcmalloc: large alloc 1212309504 bytes == 0x5580931bc000 @  0x7f5039662887 0x7f5037f58c29 0x7f5037f58d47 0x7f5037f5a7a5 0x7f4fdcde9c98 0x7f4fdcdddde2 0x558043e9ef78 0x558043f12a6d 0x558043f0d02f 0x558043e9faba 0x558043f0e108 0x558043e9f9da 0x558043f0deae 0x558043e9f9da 0x558043f122c0 0x558043f0d66e 0x558043e9faba 0x558043f0deae 0x558043f0d02f 0x558043e9faba 0x558043f0ecd4 0x558043f0d66e 0x558043e9faba 0x558043f0deae 0x558043f0d02f 0x558043e9faba 0x558043f0ecd4 0x558043e9f9da 0x558043f0deae 0x558043f0d02f 0x558043f0cd43
2022-02-23 03:08:59,135 - INFO - tf2onnx.tf_loader: Signatures found in model: [serving_default].
2022-02-23 03:08:59,135 - WARNING - tf2onnx.tf_loader: '--signature_def' not specified, using first signature: serving_default
2022-02-23 03:08:59,135 - INFO - tf2onnx.tf_loader: Output names: ['output_0']
WARNING:tensorflow:From /usr/local/lib/python3.7/dist-packages/tf2onnx/tf_loader.py:711: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.compat.v1.graph_util.extract_sub_graph`
2022-02-23 03:09:30,363 - WARNING - tensorflow: From /usr/local/lib/python3.7/dist-packages/tf2onnx/tf_loader.py:711: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.compat.v1.graph_util.extract_sub_graph`
2022-02-23 03:09:37,025 - INFO - tf2onnx: inputs: ['args_0:0', 'args_1:0', 'args_2:0']
2022-02-23 03:09:37,025 - INFO - tf2onnx: outputs: ['Identity:0']
2022-02-23 03:09:40,603 - INFO - tf2onnx.tfonnx: Using tensorflow=2.8.0, onnx=1.11.0, tf2onnx=1.10.0/bf4a22
2022-02-23 03:09:40,603 - INFO - tf2onnx.tfonnx: Using opset <onnx, 15>
2022-02-23 03:09:48,819 - INFO - tf2onnx.tf_utils: Computed 0 values for constant folding
2022-02-23 03:09:52,598 - INFO - tf2onnx.tf_utils: Computed 0 values for constant folding
2022-02-23 03:09:53,938 - VERBOSE - tf2onnx.tfonnx: Mapping TF node to ONNX node(s)
2022-02-23 03:09:53,956 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_11: BitwiseAnd] is not supported
2022-02-23 03:09:53,960 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_10: BitwiseAnd] is not supported
2022-02-23 03:09:53,964 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_9: BitwiseAnd] is not supported
2022-02-23 03:09:53,968 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_8: BitwiseAnd] is not supported
2022-02-23 03:09:53,971 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_7: BitwiseAnd] is not supported
2022-02-23 03:09:53,975 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_6: BitwiseAnd] is not supported
2022-02-23 03:09:53,978 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_5: BitwiseAnd] is not supported
2022-02-23 03:09:53,982 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_4: BitwiseAnd] is not supported
2022-02-23 03:09:53,985 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_3: BitwiseAnd] is not supported
2022-02-23 03:09:53,988 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_2: BitwiseAnd] is not supported
2022-02-23 03:09:53,991 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd_1: BitwiseAnd] is not supported
2022-02-23 03:09:53,995 - ERROR - tf2onnx.tfonnx: Tensorflow op [BitwiseAnd: BitwiseAnd] is not supported
2022-02-23 03:09:54,261 - ERROR - tf2onnx.tfonnx: Unsupported ops: Counter({'BitwiseAnd': 12})
2022-02-23 03:09:54,402 - VERBOSE - tf2onnx.tfonnx: Summay Stats:
	tensorflow ops: Counter({'Const': 1390, 'Mul': 381, 'AddV2': 336, 'Reshape': 264, 'Identity': 211, 'Einsum': 194, 'Sum': 128, 'Sub': 126, 'BroadcastTo': 69, 'Maximum': 53, 'Rsqrt': 51, 'Exp': 48, 'Cast': 36, 'Transpose': 30, 'RealDiv': 26, 'Max': 24, 'StopGradient': 24, 'Neg': 24, 'Reciprocal': 24, 'SelectV2': 14, 'NotEqual': 12, 'BitwiseAnd': 12, 'Greater': 12, 'StridedSlice': 4, 'Placeholder': 3, 'Less': 2, 'Squeeze': 2, 'Minimum': 2, 'GatherV2': 2, 'Sqrt': 2, 'ArgMax': 1, 'Conv2D': 1, 'ConcatV2': 1, 'MatMul': 1, 'IdentityN': 1})
	tensorflow attr: Counter({'dtype': 1393, 'value': 1390, 'Tidx': 223, 'N': 195, 'equation': 194, 'keep_dims': 152, 'Truncate': 36, 'to': 36, 'incompatible_shape_error': 12, 'begin_mask': 4, 'ellipsis_mask': 4, 'end_mask': 4, 'new_axis_mask': 4, 'shrink_axis_mask': 4, '_user_specified_name': 3, 'shape': 3, 'squeeze_dims': 2, 'batch_dims': 2, 'output_type': 1, 'data_format': 1, 'dilations': 1, 'explicit_paddings': 1, 'padding': 1, 'strides': 1, 'transpose_a': 1, 'transpose_b': 1, '_gradient_op_type': 1})
	onnx mapped: Counter({'Const': 1153, 'Mul': 381, 'AddV2': 336, 'Reshape': 264, 'Identity': 210, 'Einsum': 194, 'Sum': 128, 'Sub': 126, 'BroadcastTo': 69, 'Maximum': 53, 'Rsqrt': 51, 'Exp': 48, 'Cast': 36, 'Transpose': 30, 'RealDiv': 26, 'Max': 24, 'StopGradient': 24, 'Neg': 24, 'Reciprocal': 24, 'SelectV2': 14, 'NotEqual': 12, 'Greater': 12, 'StridedSlice': 4, 'Placeholder': 3, 'Sqrt': 2, 'Less': 2, 'Squeeze': 2, 'Minimum': 2, 'GatherV2': 2, 'Conv2D': 1, 'ConcatV2': 1, 'MatMul': 1, 'ArgMax': 1, 'IdentityN': 1})
	onnx unmapped: Counter({'BitwiseAnd': 12})
2022-02-23 03:09:54,405 - VERBOSE - tf2onnx.tfonnx: Mapping TF node to ONNX node(s)
2022-02-23 03:09:54,405 - ERROR - tf2onnx.tfonnx: Tensorflow op [PartitionedCall: PartitionedCall] is not supported
2022-02-23 03:09:54,405 - ERROR - tf2onnx.tfonnx: Unsupported ops: Counter({'PartitionedCall': 1})
2022-02-23 03:09:54,405 - VERBOSE - tf2onnx.tfonnx: Summay Stats:
	tensorflow ops: Counter({'Placeholder': 3, 'Identity': 2, 'PartitionedCall': 1})
	tensorflow attr: Counter({'dtype': 3, 'shape': 3, '_XlaMustCompile': 1, '_collective_manager_ids': 1, '_read_only_resource_inputs': 1, 'config': 1, 'config_proto': 1, 'executor_type': 1, 'f': 1})
	onnx mapped: Counter({'Placeholder': 3, 'Identity': 1})
	onnx unmapped: Counter({'PartitionedCall': 1})
2022-02-23 03:09:54,405 - INFO - tf2onnx.optimizer: Optimizing ONNX model
2022-02-23 03:09:54,406 - VERBOSE - tf2onnx.optimizer: Apply optimize_transpose
2022-02-23 03:09:54,407 - VERBOSE - tf2onnx.optimizer.TransposeOptimizer: no change
2022-02-23 03:09:54,407 - VERBOSE - tf2onnx.optimizer: Apply remove_redundant_upsample
2022-02-23 03:09:54,407 - VERBOSE - tf2onnx.optimizer.UpsampleOptimizer: no change
2022-02-23 03:09:54,407 - VERBOSE - tf2onnx.optimizer: Apply fold_constants
2022-02-23 03:09:54,408 - VERBOSE - tf2onnx.optimizer.ConstFoldOptimizer: no change
2022-02-23 03:09:54,408 - VERBOSE - tf2onnx.optimizer: Apply const_dequantize_optimizer
2022-02-23 03:09:54,409 - VERBOSE - tf2onnx.optimizer.ConstDequantizeOptimizer: no change
2022-02-23 03:09:54,409 - VERBOSE - tf2onnx.optimizer: Apply loop_optimizer
2022-02-23 03:09:54,410 - VERBOSE - tf2onnx.optimizer.LoopOptimizer: no change
2022-02-23 03:09:54,410 - VERBOSE - tf2onnx.optimizer: Apply merge_duplication
2022-02-23 03:09:54,411 - VERBOSE - tf2onnx.optimizer.MergeDuplicatedNodesOptimizer: no change
2022-02-23 03:09:54,411 - VERBOSE - tf2onnx.optimizer: Apply reshape_optimizer
2022-02-23 03:09:54,412 - VERBOSE - tf2onnx.optimizer.ReshapeOptimizer: no change
2022-02-23 03:09:54,412 - VERBOSE - tf2onnx.optimizer: Apply global_pool_optimizer
2022-02-23 03:09:54,413 - VERBOSE - tf2onnx.optimizer.GlobalPoolOptimizer: no change
2022-02-23 03:09:54,413 - VERBOSE - tf2onnx.optimizer: Apply q_dq_optimizer
2022-02-23 03:09:54,414 - VERBOSE - tf2onnx.optimizer.QDQOptimizer: no change
2022-02-23 03:09:54,414 - VERBOSE - tf2onnx.optimizer: Apply remove_identity
2022-02-23 03:09:54,414 - VERBOSE - tf2onnx.optimizer.IdentityOptimizer: Identity -2 (2->0)
2022-02-23 03:09:54,414 - VERBOSE - tf2onnx.optimizer: Apply remove_back_to_back
2022-02-23 03:09:54,415 - VERBOSE - tf2onnx.optimizer.BackToBackOptimizer: no change
2022-02-23 03:09:54,415 - VERBOSE - tf2onnx.optimizer: Apply einsum_optimizer
2022-02-23 03:09:54,416 - VERBOSE - tf2onnx.optimizer.EinsumOptimizer: no change
2022-02-23 03:09:54,416 - VERBOSE - tf2onnx.optimizer: Apply optimize_transpose
2022-02-23 03:09:54,417 - VERBOSE - tf2onnx.optimizer.TransposeOptimizer: no change
2022-02-23 03:09:54,417 - VERBOSE - tf2onnx.optimizer: Apply remove_redundant_upsample
2022-02-23 03:09:54,417 - VERBOSE - tf2onnx.optimizer.UpsampleOptimizer: no change
2022-02-23 03:09:54,417 - VERBOSE - tf2onnx.optimizer: Apply fold_constants
2022-02-23 03:09:54,418 - VERBOSE - tf2onnx.optimizer.ConstFoldOptimizer: no change
2022-02-23 03:09:54,418 - VERBOSE - tf2onnx.optimizer: Apply const_dequantize_optimizer
2022-02-23 03:09:54,419 - VERBOSE - tf2onnx.optimizer.ConstDequantizeOptimizer: no change
2022-02-23 03:09:54,419 - VERBOSE - tf2onnx.optimizer: Apply loop_optimizer
2022-02-23 03:09:54,419 - VERBOSE - tf2onnx.optimizer.LoopOptimizer: no change
2022-02-23 03:09:54,419 - VERBOSE - tf2onnx.optimizer: Apply merge_duplication
2022-02-23 03:09:54,420 - VERBOSE - tf2onnx.optimizer.MergeDuplicatedNodesOptimizer: no change
2022-02-23 03:09:54,420 - VERBOSE - tf2onnx.optimizer: Apply reshape_optimizer
2022-02-23 03:09:54,421 - VERBOSE - tf2onnx.optimizer.ReshapeOptimizer: no change
2022-02-23 03:09:54,421 - VERBOSE - tf2onnx.optimizer: Apply global_pool_optimizer
2022-02-23 03:09:54,423 - VERBOSE - tf2onnx.optimizer.GlobalPoolOptimizer: no change
2022-02-23 03:09:54,423 - VERBOSE - tf2onnx.optimizer: Apply q_dq_optimizer
2022-02-23 03:09:54,423 - VERBOSE - tf2onnx.optimizer.QDQOptimizer: no change
2022-02-23 03:09:54,424 - VERBOSE - tf2onnx.optimizer: Apply remove_identity
2022-02-23 03:09:54,424 - VERBOSE - tf2onnx.optimizer.IdentityOptimizer: no change
2022-02-23 03:09:54,424 - VERBOSE - tf2onnx.optimizer: Apply remove_back_to_back
2022-02-23 03:09:54,425 - VERBOSE - tf2onnx.optimizer.BackToBackOptimizer: no change
2022-02-23 03:09:54,425 - VERBOSE - tf2onnx.optimizer: Apply einsum_optimizer
2022-02-23 03:09:54,425 - VERBOSE - tf2onnx.optimizer.EinsumOptimizer: no change
2022-02-23 03:09:54,426 - INFO - tf2onnx.optimizer: After optimization: Identity -2 (2->0)
2022-02-23 03:09:54,426 - INFO - tf2onnx: 
2022-02-23 03:09:54,427 - INFO - tf2onnx: Successfully converted TensorFlow model /content/pixel_text_score to ONNX
2022-02-23 03:09:54,427 - INFO - tf2onnx: Model inputs: ['args_0', 'args_1', 'args_2']
2022-02-23 03:09:54,427 - INFO - tf2onnx: Model outputs: ['output_0']
2022-02-23 03:09:54,427 - INFO - tf2onnx: ONNX model is saved at pixel_text_score.onnx

[josephrocca]

Hey @fatcat-z I'm wondering if there is any expected ETA on support for these ops? Particularly for PartitionedCall, since I'm running into it with other models too like with this web test of JAX-MD.

I'm obviously in no position to try to rush you on this, but just curious if there is a rough estimate to help me plan :) No worries if not. Thanks!

[fatcat-z]

No, there is no ETA for these ops yet, but they are on the plate. And any contributions to this are definitely welcome!

[josephrocca]

I noticed that this has the tag "pending on user response", so I'm just responding in case this issue is stalled due to that tag. Also might as well note here that conversion of LiT runs into the same problem - lack of support for PartitionedCall (minimal replication notebook).

[josephrocca]

(Removing BitwiseAnd from title, since that's covered by this issue: #1985)

[metalMajor]

Did you find a solution for this?