Bidirectional LSTM conversion issue from PyTorch to CoreML

## 🐞Describe the bug
- Numerical difference in output of bidirectional LSTM from pytorch
- Don't know if it's a converter issue or not

## Trace
No

## To Reproduce
- Adapted LSTM test case from coremltools that reproduce the issue

```
import torch
import torch.nn as nn
import numpy as np
from six import string_types as _string_types

from coremltools import TensorType
from coremltools.converters.mil.testing_reqs import _converter
from coremltools.models import MLModel
from coremltools._deps import _IS_MACOS

def flatten_and_detach_torch_results(torch_results):
if isinstance(torch_results, (list, tuple)):
return [x.detach().numpy() for x in _flatten(torch_results)]
# Do not need to flatten
return [torch_results.detach().numpy()]

def _flatten(object):
flattened_list = []
for item in object:
if isinstance(item, (list, tuple)):
flattened_list.extend(_flatten(item))
else:
flattened_list.append(item)
return flattened_list

def generate_input_data(input_size):
if isinstance(input_size, list):
return [torch.rand(_size) for _size in input_size]
else:
return torch.rand(input_size)

def trace_model(model, input_data):
model.eval()
if isinstance(input_data, list):
input_data = tuple(input_data)
torch_model = torch.jit.trace(model, input_data)
return torch_model

def convert_to_coreml_inputs(input_description, inputs):
"""Convenience function to combine a CoreML model's input description and
set of raw inputs into the format expected by the model's predict function.
"""
flattened_inputs = _flatten(inputs)
coreml_inputs = {
str(x): inp.numpy() for x, inp in zip(input_description, flattened_inputs)
}
return coreml_inputs

def convert_to_mlmodel(model_spec, tensor_inputs, backend="nn_proto"):
def _convert_to_inputtype(inputs):
if isinstance(inputs, list):
return [_convert_to_inputtype(x) for x in inputs]
elif isinstance(inputs, tuple):
return tuple([_convert_to_inputtype(x) for x in inputs])
elif isinstance(inputs, torch.Tensor):
return TensorType(shape=inputs.shape)
else:
raise ValueError(
"Unable to parse type {} into InputType.".format(type(inputs))
)

inputs = list(_convert_to_inputtype(tensor_inputs))
proto = _converter._convert(model_spec, inputs=inputs, convert_to=backend, convert_from="torch")
return MLModel(proto, useCPUOnly=True)

def convert_and_compare(input_data, model_spec, expected_results=None, atol=1e-5, backend="nn_proto"):
"""
If expected results is not set, it will by default
be set to the flattened output of the torch model.
"""
if isinstance(model_spec, _string_types):
torch_model = torch.jit.load(model_spec)
else:
torch_model = model_spec

if not isinstance(input_data, (list, tuple)):
input_data = [input_data]

if not expected_results:
expected_results = torch_model(*input_data)
expected_results = flatten_and_detach_torch_results(expected_results)
mlmodel = convert_to_mlmodel(model_spec, input_data, backend=backend)
coreml_inputs = convert_to_coreml_inputs(mlmodel.input_description, input_data)
if _IS_MACOS:
coreml_results = mlmodel.predict(coreml_inputs)
sorted_coreml_results = [
coreml_results[key] for key in sorted(coreml_results.keys())
]

for torch_result, coreml_result in zip(expected_results, sorted_coreml_results):
np.testing.assert_equal(coreml_result.shape, torch_result.shape)
np.testing.assert_allclose(coreml_result, torch_result, atol=atol)

def run_compare_torch(
input_data, model, expected_results=None, places=5, input_as_shape=True, backend="nn_proto"
):
model.eval()
if input_as_shape:
input_data = generate_input_data(input_data)
model_spec = trace_model(model, input_data)
convert_and_compare(
input_data, model_spec, expected_results=expected_results, atol=10.0 ** -places, backend=backend
)

def _pytorch_hidden_to_coreml(x):
f, b = torch.split(x, [1] * x.shape[0], dim=0)
x = torch.cat((f, b), dim=2)
return x

def test_lstm(
input_size,
hidden_size,
num_layers,
bias,
batch_first,
dropout,
bidirectional,
backend,
):
model = nn.LSTM(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
bias=bias,
batch_first=batch_first,
dropout=dropout,
bidirectional=bidirectional,
)
SEQUENCE_LENGTH = 1
BATCH_SIZE = 1

num_directions = int(bidirectional) + 1

# (seq_len, batch, input_size)
if batch_first:
_input = torch.rand(BATCH_SIZE, SEQUENCE_LENGTH, input_size)
else:
_input = torch.randn(SEQUENCE_LENGTH, BATCH_SIZE, input_size)

h0 = torch.randn(num_layers * num_directions, BATCH_SIZE, hidden_size)
c0 = torch.randn(num_layers * num_directions, BATCH_SIZE, hidden_size)

inputs = (_input, (h0, c0))
expected_results = model(*inputs)
# Need to do some output reshaping if bidirectional
if bidirectional:
ex_hn = _pytorch_hidden_to_coreml(expected_results[1][0])
ex_cn = _pytorch_hidden_to_coreml(expected_results[1][1])
expected_results = (expected_results[0], (ex_hn, ex_cn))
run_compare_torch(inputs, model, expected_results, input_as_shape=False, backend=backend)

test_lstm(2, 128, 1, True, False, 0, True, 'nn_proto')
```

Results:

```
AssertionError:
Not equal to tolerance rtol=1e-07, atol=1e-05

Mismatch: 50%
Max absolute difference: 0.9107815
Max relative difference: 36.60195
x: array([[[ 0.159638, -0.153935, 0.07789 , 0.005968, 0.144887,
-0.186706, 0.484189, -0.341685, -0.005737, -0.285721,
0.446474, 0.461001, -0.303098, -0.127876, 0.119229,...
y: array([[[ 0.159638, -0.153935, 0.07789 , 0.005968, 0.144887,
-0.186706, 0.484189, -0.341685, -0.005737, -0.285721,
0.446474, 0.461001, -0.303098, -0.127876, 0.119229,...
```

## System environment (please complete the following information):
- coremltools version (e.g., 3.0b5): 4.0b1 & 4.0b2
- OS (e.g., MacOS, Linux): MacOS
- macOS version (if applicable): 10.15.5
- XCode version (if applicable): n/a
- How you install python (anaconda, virtualenv, system): Anaconda
- python version (e.g. 3.7): 3.6.9
- any other relevant information:

## Additional context

Start with the provided test_lstm reproduction and the _converter._convert call using convert_from="torch" and convert_to="nn_proto". Compare the bidirectional LSTM output and hidden-state reshaping against the PyTorch results, then add or update a focused converter test that passes the stated tolerance on macOS.

Written by the indexing model from the issue text.