Repository: mlflow/mlflow

Issue:
[BUG] Test Failures on master `test_onnx_model_export.py`
Thank you for submitting an issue. Please refer to our [issue policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md) for additional information about bug reports. For help with debugging your code, please refer to [Stack Overflow](https://stackoverflow.com/questions/tagged/mlflow).

Failed test cases in `tests/onnx/test_onnx_model_export.py`  and linting errors in `test_tensorflow2_autolog`. Possibly due to newly released tensorflow version. https://pypi.org/project/tensorflow/ 

### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [ ] Yes. I can contribute a fix for this bug independently.
- [ ] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**:
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**:
- **MLflow installed from (source or binary)**:
- **MLflow version (run ``mlflow --version``)**:
- **Python version**:
- **npm version, if running the dev UI**:
- **Exact command to reproduce**:



### Describe the problem
Describe the problem clearly here. Include descriptions of the expected behavior and the actual behavior.

### Code to reproduce issue
Provide a reproducible test case that is the bare minimum necessary to generate the problem.

### Other info / logs
Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached.


### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: Local serving, model deployment tools, spark UDFs
- [ ] `area/server-infra`: MLflow server, JavaScript dev server
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, JavaScript, plotting
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:
<img width="922" alt="Screen Shot 2020-12-14 at 2 36 58 PM" src="https://user-images.githubusercontent.com/73549313/102146030-29ba9700-3e1d-11eb-9c89-f29f0cd0be9b.png">


Resolving patch:
diff --git a/mlflow/spark.py b/mlflow/spark.py
--- a/mlflow/spark.py
+++ b/mlflow/spark.py
@@ -631,8 +631,8 @@ def predict(self, pandas_df):
 @autologging_integration(FLAVOR_NAME)
 def autolog(disable=False):  # pylint: disable=unused-argument
     """
-    Enables (or disables) and configures logging of Spark datasource paths, versions (if applicable), and formats
-    when they are read. This method is not threadsafe and assumes a
+    Enables (or disables) and configures logging of Spark datasource paths, versions
+    (if applicable), and formats when they are read. This method is not threadsafe and assumes a
     `SparkSession
     <https://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.SparkSession>`_
     already exists with the


Test patch:
diff --git a/tests/onnx/test_onnx_model_export.py b/tests/onnx/test_onnx_model_export.py
--- a/tests/onnx/test_onnx_model_export.py
+++ b/tests/onnx/test_onnx_model_export.py
@@ -3,14 +3,17 @@
 import pytest
 from unittest import mock
 
-from keras.models import Sequential
-from keras.layers import Dense
+import onnx
+import torch
+import torch.nn as nn
+import torch.onnx
+from torch.utils.data import DataLoader
 import sklearn.datasets as datasets
 import pandas as pd
 import numpy as np
 import yaml
 
-import tensorflow.compat.v1 as tf  # pylint: disable=import-error
+import mlflow.onnx
 import mlflow.pyfunc.scoring_server as pyfunc_scoring_server
 from mlflow import pyfunc
 from mlflow.models import infer_signature, Model
@@ -22,6 +25,10 @@
 from mlflow.utils.model_utils import _get_flavor_configuration
 from mlflow.tracking._model_registry import DEFAULT_AWAIT_MAX_SLEEP_SECONDS
 
+
+TEST_DIR = "tests"
+TEST_ONNX_RESOURCES_DIR = os.path.join(TEST_DIR, "resources", "onnx")
+
 pytestmark = pytest.mark.skipif(
     (sys.version_info < (3, 6)), reason="Tests require Python 3 to run!"
 )
@@ -39,94 +46,91 @@ def data():
 
 
 @pytest.fixture(scope="module")
-def model(data):
+def dataset(data):
     x, y = data
-    model = Sequential()
-    model.add(Dense(3, input_dim=4))
-    model.add(Dense(1))
-    model.compile(loss="mean_squared_error", optimizer="SGD")
-    model.fit(x, y)
-    return model
+    dataset = [(xi.astype(np.float32), yi.astype(np.float32)) for xi, yi in zip(x.values, y.values)]
+    return dataset
 
 
 @pytest.fixture(scope="module")
-def onnx_model(model):
-    import onnxmltools
-
-    return onnxmltools.convert_keras(model)
+def sample_input(dataset):
+    dataloader = DataLoader(dataset, batch_size=5, num_workers=1,)
+    # Load a batch from the data loader and return the samples
+    x, _ = next(iter(dataloader))
+    return x
 
 
 @pytest.fixture(scope="module")
-def sklearn_model(data):
-    from sklearn.linear_model import LogisticRegression
-
-    x, y = data
-    model = LogisticRegression()
-    model.fit(x, y)
-    return model
-
-
-@pytest.fixture(scope="module")
-def onnx_sklearn_model(sklearn_model):
-    import onnxmltools
-    from skl2onnx.common.data_types import FloatTensorType
-
-    initial_type = [("float_input", FloatTensorType([None, 4]))]
-    onx = onnxmltools.convert_sklearn(sklearn_model, initial_types=initial_type)
-    return onx
+def model(dataset):
+    model = nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1),)
+    criterion = nn.MSELoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
+    batch_size = 16
+    num_workers = 4
+    dataloader = DataLoader(
+        dataset, batch_size=batch_size, num_workers=num_workers, shuffle=True, drop_last=False
+    )
 
+    model.train()
+    for _ in range(5):
+        for batch in dataloader:
+            optimizer.zero_grad()
+            batch_size = batch[0].shape[0]
+            y_pred = model(batch[0]).squeeze(dim=1)
+            loss = criterion(y_pred, batch[1])
+            loss.backward()
+            optimizer.step()
 
-@pytest.fixture(scope="module")
-def predicted(model, data):
-    return model.predict(data[0])
+    return model
 
 
-@pytest.fixture(scope="module")
-def tf_model_multiple_inputs_float64():
-    graph = tf.Graph()
-    with graph.as_default():
-        t_in1 = tf.placeholder(tf.float64, 10, name="first_input")
-        t_in2 = tf.placeholder(tf.float64, 10, name="second_input")
-        t_out = tf.multiply(t_in1, t_in2)
-        tf.identity(t_out, name="output")
-    return graph
+@pytest.fixture
+def onnx_model(model, sample_input, tmpdir):
+    model_path = os.path.join(str(tmpdir), "torch_onnx")
+    dynamic_axes = {"input": {0: "batch"}}
+    torch.onnx.export(
+        model, sample_input, model_path, dynamic_axes=dynamic_axes, input_names=["input"]
+    )
+    return onnx.load(model_path)
 
 
 @pytest.fixture(scope="module")
-def tf_model_multiple_inputs_float32():
-    graph = tf.Graph()
-    with graph.as_default():
-        t_in1 = tf.placeholder(tf.float32, 10, name="first_input")
-        t_in2 = tf.placeholder(tf.float32, 10, name="second_input")
-        t_out = tf.multiply(t_in1, t_in2)
-        tf.identity(t_out, name="output")
-    return graph
+def onnx_sklearn_model():
+    """
+    A scikit-learn model in ONNX format that is used to test the behavior
+    of ONNX models that return outputs in list format. For reference, see
+    `test_pyfunc_predict_supports_models_with_list_outputs`.
+    """
+    model_path = os.path.join(TEST_ONNX_RESOURCES_DIR, "sklearn_model.onnx")
+    return onnx.load(model_path)
 
 
 @pytest.fixture(scope="module")
-def onnx_model_multiple_inputs_float64(tf_model_multiple_inputs_float64):
-    import tf2onnx
-
-    sess = tf.Session(graph=tf_model_multiple_inputs_float64)
-
-    onnx_graph = tf2onnx.tfonnx.process_tf_graph(
-        sess.graph, input_names=["first_input:0", "second_input:0"], output_names=["output:0"]
+def predicted(model, dataset):
+    batch_size = 16
+    num_workers = 4
+    dataloader = DataLoader(
+        dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False, drop_last=False
     )
-    model_proto = onnx_graph.make_model("test")
-    return model_proto
+    predictions = np.zeros((len(dataloader.sampler),))
+    model.eval()
+    with torch.no_grad():
+        for i, batch in enumerate(dataloader):
+            y_preds = model(batch[0]).squeeze(dim=1).numpy()
+            predictions[i * batch_size : (i + 1) * batch_size] = y_preds
+    return predictions
 
 
 @pytest.fixture(scope="module")
-def onnx_model_multiple_inputs_float32(tf_model_multiple_inputs_float32):
-    import tf2onnx
+def onnx_model_multiple_inputs_float64():
+    model_path = os.path.join(TEST_ONNX_RESOURCES_DIR, "tf_model_multiple_inputs_float64.onnx")
+    return onnx.load(model_path)
 
-    sess = tf.Session(graph=tf_model_multiple_inputs_float32)
 
-    onnx_graph = tf2onnx.tfonnx.process_tf_graph(
-        sess.graph, input_names=["first_input:0", "second_input:0"], output_names=["output:0"]
-    )
-    model_proto = onnx_graph.make_model("test")
-    return model_proto
+@pytest.fixture(scope="module")
+def onnx_model_multiple_inputs_float32():
+    model_path = os.path.join(TEST_ONNX_RESOURCES_DIR, "tf_model_multiple_inputs_float32.onnx")
+    return onnx.load(model_path)
 
 
 @pytest.fixture(scope="module")
@@ -152,17 +156,13 @@ def model_path(tmpdir):
 def onnx_custom_env(tmpdir):
     conda_env = os.path.join(str(tmpdir), "conda_env.yml")
     _mlflow_conda_env(
-        conda_env,
-        additional_conda_deps=["pytest", "keras"],
-        additional_pip_deps=["onnx", "onnxmltools"],
+        conda_env, additional_conda_deps=["pytest", "torch"], additional_pip_deps=["onnx"],
     )
     return conda_env
 
 
 @pytest.mark.large
 def test_cast_float64_to_float32():
-    import mlflow.onnx
-
     df = pd.DataFrame([[1.0, 2.1], [True, False]], columns=["col1", "col2"])
     df["col1"] = df["col1"].astype(np.float64)
     df["col2"] = df["col2"].astype(np.bool)
@@ -170,14 +170,9 @@ def test_cast_float64_to_float32():
     assert df2["col1"].dtype == np.float32 and df2["col2"].dtype == np.bool
 
 
-# TODO: Use the default conda environment once MLflow's Travis build supports the onnxruntime
-# library
 @pytest.mark.large
-def test_model_save_load(onnx_model, model_path, onnx_custom_env):
-    import onnx
-    import mlflow.onnx
-
-    mlflow.onnx.save_model(onnx_model, model_path, conda_env=onnx_custom_env)
+def test_model_save_load(onnx_model, model_path):
+    mlflow.onnx.save_model(onnx_model, model_path)
 
     # Loading ONNX model
     onnx.checker.check_model = mock.Mock()
@@ -187,8 +182,6 @@ def test_model_save_load(onnx_model, model_path, onnx_custom_env):
 
 @pytest.mark.large
 def test_signature_and_examples_are_saved_correctly(onnx_model, data, onnx_custom_env):
-    import mlflow.onnx
-
     model = onnx_model
     signature_ = infer_signature(*data)
     example_ = data[0].head(3)
@@ -211,17 +204,14 @@ def test_signature_and_examples_are_saved_correctly(onnx_model, data, onnx_custo
                     assert all((_read_example(mlflow_model, path) == example).all())
 
 
-# TODO: Mark this as large once MLflow's Travis build supports the onnxruntime library
-@pytest.mark.release
+@pytest.mark.large
 def test_model_save_load_evaluate_pyfunc_format(onnx_model, model_path, data, predicted):
-    import mlflow.onnx
-
     x = data[0]
     mlflow.onnx.save_model(onnx_model, model_path)
 
     # Loading pyfunc model
     pyfunc_loaded = mlflow.pyfunc.load_pyfunc(model_path)
-    assert np.allclose(pyfunc_loaded.predict(x).values, predicted, rtol=1e-05, atol=1e-05)
+    assert np.allclose(pyfunc_loaded.predict(x).values.flatten(), predicted, rtol=1e-05, atol=1e-05)
 
     # pyfunc serve
     scoring_response = pyfunc_serve_and_score_model(
@@ -230,25 +220,18 @@ def test_model_save_load_evaluate_pyfunc_format(onnx_model, model_path, data, pr
         content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON_SPLIT_ORIENTED,
     )
     assert np.allclose(
-        pd.read_json(scoring_response.content, orient="records").values.astype(np.float32),
+        pd.read_json(scoring_response.content, orient="records")
+        .values.flatten()
+        .astype(np.float32),
         predicted,
         rtol=1e-05,
         atol=1e-05,
     )
 
 
-# TODO: Use the default conda environment once MLflow's Travis build supports the onnxruntime
-# library
 @pytest.mark.large
-def test_model_save_load_multiple_inputs(
-    onnx_model_multiple_inputs_float64, model_path, onnx_custom_env
-):
-    import onnx
-    import mlflow.onnx
-
-    mlflow.onnx.save_model(
-        onnx_model_multiple_inputs_float64, model_path, conda_env=onnx_custom_env
-    )
+def test_model_save_load_multiple_inputs(onnx_model_multiple_inputs_float64, model_path):
+    mlflow.onnx.save_model(onnx_model_multiple_inputs_float64, model_path)
 
     # Loading ONNX model
     onnx.checker.check_model = mock.Mock()
@@ -256,13 +239,10 @@ def test_model_save_load_multiple_inputs(
     assert onnx.checker.check_model.called
 
 
-# TODO: Mark this as large once MLflow's Travis build supports the onnxruntime library
-@pytest.mark.release
+@pytest.mark.large
 def test_model_save_load_evaluate_pyfunc_format_multiple_inputs(
     onnx_model_multiple_inputs_float64, data_multiple_inputs, predicted_multiple_inputs, model_path
 ):
-    import mlflow.onnx
-
     mlflow.onnx.save_model(onnx_model_multiple_inputs_float64, model_path)
 
     # Loading pyfunc model
@@ -290,8 +270,7 @@ def test_model_save_load_evaluate_pyfunc_format_multiple_inputs(
 
 # TODO: Remove test, along with explicit casting, when https://github.com/mlflow/mlflow/issues/1286
 # is fixed.
-# TODO: Mark this as large once MLflow's Travis build supports the onnxruntime library
-@pytest.mark.release
+@pytest.mark.large
 def test_pyfunc_representation_of_float32_model_casts_and_evalutes_float64_inputs(
     onnx_model_multiple_inputs_float32, model_path, data_multiple_inputs, predicted_multiple_inputs
 ):
@@ -304,8 +283,6 @@ def test_pyfunc_representation_of_float32_model_casts_and_evalutes_float64_input
     precision (e.g., 32-bit floats may be converted to 64-bit floats when persisting a
     DataFrame as JSON).
     """
-    import mlflow.onnx
-
     mlflow.onnx.save_model(onnx_model_multiple_inputs_float32, model_path)
 
     # Loading pyfunc model
@@ -317,28 +294,19 @@ def test_pyfunc_representation_of_float32_model_casts_and_evalutes_float64_input
         atol=1e-05,
     )
 
-    with pytest.raises(RuntimeError):
+    with pytest.raises(Exception, match="Unexpected input data type"):
         pyfunc_loaded.predict(data_multiple_inputs.astype("int32"))
 
 
-# TODO: Use the default conda environment once MLflow's Travis build supports the onnxruntime
-# library
 @pytest.mark.large
-def test_model_log(onnx_model, onnx_custom_env):
-    # pylint: disable=unused-argument
-
-    import onnx
-    import mlflow.onnx
-
+def test_model_log(onnx_model):
     # should_start_run tests whether or not calling log_model() automatically starts a run.
     for should_start_run in [False, True]:
         try:
             if should_start_run:
                 mlflow.start_run()
             artifact_path = "onnx_model"
-            mlflow.onnx.log_model(
-                onnx_model=onnx_model, artifact_path=artifact_path, conda_env=onnx_custom_env
-            )
+            mlflow.onnx.log_model(onnx_model=onnx_model, artifact_path=artifact_path)
             model_uri = "runs:/{run_id}/{artifact_path}".format(
                 run_id=mlflow.active_run().info.run_id, artifact_path=artifact_path
             )
@@ -352,8 +320,6 @@ def test_model_log(onnx_model, onnx_custom_env):
 
 
 def test_log_model_calls_register_model(onnx_model, onnx_custom_env):
-    import mlflow.onnx
-
     artifact_path = "model"
     register_model_patch = mock.patch("mlflow.register_model")
     with mlflow.start_run(), register_model_patch:
@@ -372,8 +338,6 @@ def test_log_model_calls_register_model(onnx_model, onnx_custom_env):
 
 
 def test_log_model_no_registered_model_name(onnx_model, onnx_custom_env):
-    import mlflow.onnx
-
     artifact_path = "model"
     register_model_patch = mock.patch("mlflow.register_model")
     with mlflow.start_run(), register_model_patch:
@@ -383,12 +347,10 @@ def test_log_model_no_registered_model_name(onnx_model, onnx_custom_env):
         mlflow.register_model.assert_not_called()
 
 
-# TODO: Mark this as large once MLflow's Travis build supports the onnxruntime library
-@pytest.mark.release
+@pytest.mark.large
 def test_model_log_evaluate_pyfunc_format(onnx_model, data, predicted):
-    import mlflow.onnx
-
     x = data[0]
+
     # should_start_run tests whether or not calling log_model() automatically starts a run.
     for should_start_run in [False, True]:
         try:
@@ -402,7 +364,9 @@ def test_model_log_evaluate_pyfunc_format(onnx_model, data, predicted):
 
             # Loading pyfunc model
             pyfunc_loaded = mlflow.pyfunc.load_pyfunc(model_uri=model_uri)
-            assert np.allclose(pyfunc_loaded.predict(x).values, predicted, rtol=1e-05, atol=1e-05)
+            assert np.allclose(
+                pyfunc_loaded.predict(x).values.flatten(), predicted, rtol=1e-05, atol=1e-05
+            )
         finally:
             mlflow.end_run()
 
@@ -411,8 +375,6 @@ def test_model_log_evaluate_pyfunc_format(onnx_model, data, predicted):
 def test_model_save_persists_specified_conda_env_in_mlflow_model_directory(
     onnx_model, model_path, onnx_custom_env
 ):
-    import mlflow.onnx
-
     mlflow.onnx.save_model(onnx_model=onnx_model, path=model_path, conda_env=onnx_custom_env)
     pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
     saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
@@ -426,11 +388,8 @@ def test_model_save_persists_specified_conda_env_in_mlflow_model_directory(
     assert saved_conda_env_parsed == onnx_custom_env_parsed
 
 
-# TODO: Mark this as large once MLflow's Travis build supports the onnxruntime library
-@pytest.mark.release
+@pytest.mark.large
 def test_model_save_accepts_conda_env_as_dict(onnx_model, model_path):
-    import mlflow.onnx
-
     conda_env = dict(mlflow.onnx.get_default_conda_env())
     conda_env["dependencies"].append("pytest")
     mlflow.onnx.save_model(onnx_model=onnx_model, path=model_path, conda_env=conda_env)
@@ -448,8 +407,6 @@ def test_model_save_accepts_conda_env_as_dict(onnx_model, model_path):
 def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(
     onnx_model, onnx_custom_env
 ):
-    import mlflow.onnx
-
     artifact_path = "model"
     with mlflow.start_run():
         mlflow.onnx.log_model(
@@ -473,13 +430,10 @@ def test_model_log_persists_specified_conda_env_in_mlflow_model_directory(
     assert saved_conda_env_parsed == onnx_custom_env_parsed
 
 
-# TODO: Mark this as large once MLflow's Travis build supports the onnxruntime library
-@pytest.mark.release
+@pytest.mark.large
 def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
     onnx_model, model_path
 ):
-    import mlflow.onnx
-
     mlflow.onnx.save_model(onnx_model=onnx_model, path=model_path, conda_env=None)
     pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME)
     conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV])
@@ -489,13 +443,10 @@ def test_model_save_without_specified_conda_env_uses_default_env_with_expected_d
     assert conda_env == mlflow.onnx.get_default_conda_env()
 
 
-# TODO: Mark this as large once MLflow's Travis build supports the onnxruntime library
-@pytest.mark.release
+@pytest.mark.large
 def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(
     onnx_model,
 ):
-    import mlflow.onnx
-
     artifact_path = "model"
     with mlflow.start_run():
         mlflow.onnx.log_model(onnx_model=onnx_model, artifact_path=artifact_path, conda_env=None)
@@ -513,8 +464,7 @@ def test_model_log_without_specified_conda_env_uses_default_env_with_expected_de
     assert conda_env == mlflow.onnx.get_default_conda_env()
 
 
-# TODO: Mark this as large once MLflow's Travis build supports the onnxruntime library
-@pytest.mark.release
+@pytest.mark.large
 def test_pyfunc_predict_supports_models_with_list_outputs(onnx_sklearn_model, model_path, data):
     """
     https://github.com/mlflow/mlflow/issues/2499
@@ -522,8 +472,6 @@ def test_pyfunc_predict_supports_models_with_list_outputs(onnx_sklearn_model, mo
     The issue resulted in an error because MLflow assumed it would be a numpy array. Therefore,
     the this test validates the service does not receive that error when using such a model.
     """
-    import mlflow.onnx
-
     x = data[0]
     mlflow.onnx.save_model(onnx_sklearn_model, model_path)
     wrapper = mlflow.pyfunc.load_model(model_path)
diff --git a/tests/resources/onnx/generate_onnx_models.py b/tests/resources/onnx/generate_onnx_models.py
new file mode 100644
--- /dev/null
+++ b/tests/resources/onnx/generate_onnx_models.py
@@ -0,0 +1,81 @@
+# pylint: disable=import-error
+
+"""
+Generates the following test resources:
+
+    - tf_model_multiple_inputs_float32.onnx
+    - tf_model_multiple_inputs_float64.onnx
+    - sklearn_model.onnx
+
+Usage: python generate_onnx_models.py
+"""
+
+
+import numpy as np
+import onnx
+import onnxmltools
+import pandas as pd
+import sklearn.datasets as datasets
+import tensorflow.compat.v1 as tf
+import tf2onnx
+from skl2onnx.common.data_types import FloatTensorType
+from sklearn.linear_model import LogisticRegression
+
+tf.disable_v2_behavior()
+
+
+def generate_tf_onnx_model_multiple_inputs_float64():
+    graph = tf.Graph()
+    with graph.as_default():
+        t_in1 = tf.placeholder(tf.float64, 10, name="first_input")
+        t_in2 = tf.placeholder(tf.float64, 10, name="second_input")
+        t_out = tf.multiply(t_in1, t_in2)
+        tf.identity(t_out, name="output")
+
+    sess = tf.Session(graph=graph)
+
+    onnx_graph = tf2onnx.tfonnx.process_tf_graph(
+        sess.graph, input_names=["first_input:0", "second_input:0"], output_names=["output:0"]
+    )
+    model_proto = onnx_graph.make_model("test")
+
+    onnx.save_model(model_proto, "tf_model_multiple_inputs_float64.onnx")
+
+
+def generate_tf_onnx_model_multiple_inputs_float32():
+    graph = tf.Graph()
+    with graph.as_default():
+        t_in1 = tf.placeholder(tf.float32, 10, name="first_input")
+        t_in2 = tf.placeholder(tf.float32, 10, name="second_input")
+        t_out = tf.multiply(t_in1, t_in2)
+        tf.identity(t_out, name="output")
+
+    sess = tf.Session(graph=graph)
+
+    onnx_graph = tf2onnx.tfonnx.process_tf_graph(
+        sess.graph, input_names=["first_input:0", "second_input:0"], output_names=["output:0"]
+    )
+    model_proto = onnx_graph.make_model("test")
+
+    onnx.save_model(model_proto, "tf_model_multiple_inputs_float32.onnx")
+
+
+def generate_sklearn_onnx_model():
+    iris = datasets.load_iris()
+    data = pd.DataFrame(
+        data=np.c_[iris["data"], iris["target"]], columns=iris["feature_names"] + ["target"]
+    )
+    y = data["target"]
+    x = data.drop("target", axis=1)
+
+    model = LogisticRegression()
+    model.fit(x, y)
+
+    initial_type = [("float_input", FloatTensorType([None, 4]))]
+    onx = onnxmltools.convert_sklearn(model, initial_types=initial_type)
+    onnx.save_model(onx, "sklearn_model.onnx")
+
+
+generate_tf_onnx_model_multiple_inputs_float32()
+generate_tf_onnx_model_multiple_inputs_float64()
+generate_sklearn_onnx_model()
diff --git a/tests/resources/onnx/sklearn_model.onnx b/tests/resources/onnx/sklearn_model.onnx
new file mode 100644
Binary files /dev/null and b/tests/resources/onnx/sklearn_model.onnx differ
diff --git a/tests/resources/onnx/tf_model_multiple_inputs_float32.onnx b/tests/resources/onnx/tf_model_multiple_inputs_float32.onnx
new file mode 100644
Binary files /dev/null and b/tests/resources/onnx/tf_model_multiple_inputs_float32.onnx differ
diff --git a/tests/resources/onnx/tf_model_multiple_inputs_float64.onnx b/tests/resources/onnx/tf_model_multiple_inputs_float64.onnx
new file mode 100644
Binary files /dev/null and b/tests/resources/onnx/tf_model_multiple_inputs_float64.onnx differ
diff --git a/tests/tensorflow_autolog/test_tensorflow2_autolog.py b/tests/tensorflow_autolog/test_tensorflow2_autolog.py
--- a/tests/tensorflow_autolog/test_tensorflow2_autolog.py
+++ b/tests/tensorflow_autolog/test_tensorflow2_autolog.py
@@ -6,7 +6,7 @@
 import numpy as np
 import pandas as pd
 import tensorflow as tf
-from tensorflow.python.keras import layers  # pylint: disable=import-error
+from tensorflow.keras import layers
 
 import mlflow
 import mlflow.tensorflow
diff --git a/tests/tensorflow_autolog/test_tensorflow_autolog.py b/tests/tensorflow_autolog/test_tensorflow_autolog.py
--- a/tests/tensorflow_autolog/test_tensorflow_autolog.py
+++ b/tests/tensorflow_autolog/test_tensorflow_autolog.py
@@ -8,7 +8,7 @@
 import numpy as np
 import pandas as pd
 import tensorflow as tf
-from tensorflow.python.keras import layers  # pylint: disable=import-error
+from tensorflow.keras import layers
 
 import mlflow
 import mlflow.tensorflow