Cant convert GRU layer

[ChouWeizhong]

🐞Describe the bug

Hi, Thank you for looking my bug report.
I'm working on Pytorch model development, and I need convert my model to CoreML
So the onlny way I know is Pytorch->Onnx->CoreML(if there are other good way please tell me)
The convert works well before I add GRU layer into my model.
Then after I add GRU: the Pytorch->Onnx works well. But onnx model cant convet to CoreML.

Trace

1/13: Converting Node Type Transpose
2/13: Converting Node Type Slice
3/13: Converting Node Type GRU
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
~/Desktop/project/location/Auuo_CoreML_models/mlvirtualenv/pythonenv/lib/python3.7/site-packages/numpy/lib/shape_base.py in split(ary, indices_or_sections, axis)
    864     try:
--> 865         len(indices_or_sections)
    866     except TypeError:

TypeError: object of type 'int' has no len()

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
<ipython-input-14-07e09e981ff3> in <module>
      2 onnx_model_path = f"out/{export_model_name}.onnx"
      3 onnx_model = onnx.load(onnx_model_path)
----> 4 cml_model= convert(onnx_model,minimum_ios_deployment_target='13')
      5 cml_model.save(f"out/{export_model_name}.mlmodel")

~/Desktop/project/location/Auuo_CoreML_models/mlvirtualenv/pythonenv/lib/python3.7/site-packages/onnx_coreml/converter.py in convert(model, mode, image_input_names, preprocessing_args, image_output_names, deprocessing_args, class_labels, predicted_feature_name, add_custom_layers, custom_conversion_functions, onnx_coreml_input_shape_map, minimum_ios_deployment_target)
    624         print("%d/%d: Converting Node Type %s" %(i+1, len(graph.nodes), node.op_type))
    625         if disable_coreml_rank5_mapping:
--> 626             _convert_node_nd(builder, node, graph, err)
    627         else:
    628             _add_const_inputs_if_required(builder, node, graph, err)

~/Desktop/project/location/Auuo_CoreML_models/mlvirtualenv/pythonenv/lib/python3.7/site-packages/onnx_coreml/_operators_nd.py in _convert_node_nd(builder, node, graph, err)
   2385 def _convert_node_nd(builder, node, graph, err):  # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None
   2386     converter_fn = _get_node_converter_fn(builder, node, err)
-> 2387     return converter_fn(builder, node, graph, err)
   2388 

~/Desktop/project/location/Auuo_CoreML_models/mlvirtualenv/pythonenv/lib/python3.7/site-packages/onnx_coreml/_operators_nd.py in _convert_gru(builder, node, graph, err)
    843 
    844     # Get weights for forward direction
--> 845     W_x, W_h, b = get_weights(W, W_name, R, R_name, B)
    846 
    847     # shape of input

~/Desktop/project/location/Auuo_CoreML_models/mlvirtualenv/pythonenv/lib/python3.7/site-packages/onnx_coreml/_operators_nd.py in get_weights(W, W_name, R, R_name, B)
    784                                     "Weight tensor: {} not found in the graph initializer".format(R_name))
    785 
--> 786         W_z, W_r, W_h = np.split(np.squeeze(W), 3)  #type: ignore
    787         R_z, R_r, R_h = np.split(np.squeeze(R), 3)  #type: ignore
    788 

<__array_function__ internals> in split(*args, **kwargs)

~/Desktop/project/location/Auuo_CoreML_models/mlvirtualenv/pythonenv/lib/python3.7/site-packages/numpy/lib/shape_base.py in split(ary, indices_or_sections, axis)
    869         if N % sections:
    870             raise ValueError(
--> 871                 'array split does not result in an equal division')
    872     return array_split(ary, indices_or_sections, axis)
    873 

ValueError: array split does not result in an equal division

To Reproduce

• If a python script can reproduce the error, please paste the code snippet

from onnx_coreml import convert
import onnx
# convert to CoreML
onnx_model_path = f"out/{export_model_name}.onnx"
onnx_model = onnx.load(onnx_model_path)
cml_model= convert(onnx_model,minimum_ios_deployment_target='13')
cml_model.save(f"out/{export_model_name}.mlmodel")

• If applicable, please attach ONNX model
  • If model cannot be shared publicly, please attach it via filing a bug report at https://developer.apple.com/bug-reporting/

I don't know how to up load to apple bug-reporting
I will share my model demo code latter.

System environment (please complete the following information):

• coremltools version : 3.3
• onnx-coreml version : 1.2
  onnxruntime 1.1.1
  onnx 1.6.0
• OS (e.g., MacOS, Linux): MacOS 10.15.4
• macOS version (if applicable): 10.15.4
• How you install python (anaconda, virtualenv, system): virtualenv
• python version (e.g. 3.7): Python 3.7.3
• any other relevant information:

Additional context

Add any other context about the problem here.

[bhushan23]

Thanks @ChouWeizhong
separating gates is incorrect (here: https://github.com/onnx/onnx-coreml/blob/master/onnx_coreml/_operators_nd.py#L786) It should have been splitting on 2nd axis always instead of squeezing which will split differently depending on input.