[BUG] infer_signature doesn't work with pandas dataframe when logging tensorflow models

[VictorAtIfInsurance]

Issues Policy acknowledgement

• [x] I have read and agree to submit bug reports in accordance with the issues policy

Where did you encounter this bug?

Databricks

Willingness to contribute

No. I cannot contribute a bug fix at this time.

MLflow version

mlflow==2.9.2

System information

Databricks:

• 13.3 LTS (includes Apache Spark 3.4.1, Scala 2.12) runtime
• MLflow and tensorflow installed inside the cluster

Locally:

• Inside a Debian GNU/Linux 11 (bullseye) docker container (built from mcr.microsoft.com/vscode/devcontainers/python:0-3.10-bullseye)
• MLflow and tensorflow installed inside container

Describe the problem

When trying to log a tensorflow model with the signature inferred from a pandas dataframe it gets the wrong signature input schema. The signature for the tensorflow.log_model method requires that both the inputs and outputs of the signature to be of type TensorSpec, but the inferred signature from the dataframe returns native python types.

Inferring the signature from a numpy array seems to work as expected though.

I get the same behaviour in databricks as locally.

Tracking information

REPLACE_ME

Code to reproduce issue

from tensorflow.keras.models import Sequential 
from tensorflow.keras.layers import Dense
import tensorflow as tf
import pandas as pd
from mlflow.models import infer_signature
import mlflow
import numpy as np

with mlflow.start_run():
    data = pd.DataFrame({"inputs": [0., 0., 0.]})
    model = Sequential()
    model.add(Dense(1, input_shape=(data.shape[-1],)))
    model.compile(loss="mse", optimizer="adam")
    signature = infer_signature(data, model.predict(data))
    print(signature)
    mlflow.tensorflow.log_model(
        model,
        artifact_path="model",
        registered_model_name="model",
        signature=signature,
    )

Stack trace

signature=inputs: 
  ['inputs': double]
outputs: 
  [Tensor('float32', (-1, 1))]
params: 
  None

Traceback (most recent call last):
  File "/workspaces/test-model-signature/register_model.py", line 17, in <module>
    mlflow.tensorflow.log_model(
  File "/usr/local/lib/python3.10/site-packages/mlflow/tensorflow/__init__.py", line 215, in log_model
    return Model.log(
  File "/usr/local/lib/python3.10/site-packages/mlflow/models/model.py", line 619, in log
    flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/mlflow/tensorflow/__init__.py", line 369, in save_model
    raise MlflowException(
mlflow.exceptions.MlflowException: All fields in the model signature's input schema must be of type TensorSpec.

Other info / logs

# This is a working example.

from tensorflow.keras.models import Sequential 
from tensorflow.keras.layers import Dense
import tensorflow as tf
import pandas as pd
from mlflow.models import infer_signature
import mlflow
import numpy as np

with mlflow.start_run():
    data = pd.DataFrame({"inputs": [0., 0., 0.]})
    # Convert the dataframe to a numpy array infers the correct signature that tensorflow.log_model expects.
    data = np.ndarray(shape=(1, 3), dtype=float, buffer=data.values)
    model = Sequential()
    model.add(Dense(1, input_shape=(data.shape[-1],)))
    model.compile(loss="mse", optimizer="adam")
    signature = infer_signature(data, model.predict(data))
    print(signature)
    mlflow.tensorflow.log_model(
        model,
        artifact_path="model",
        registered_model_name="model",
        signature=signature,
    )

What component(s) does this bug affect?

• [ ] area/artifacts: Artifact stores and artifact logging
• [ ] area/build: Build and test infrastructure for MLflow
• [ ] area/deployments: MLflow Deployments client APIs, server, and third-party Deployments integrations
• [ ] area/docs: MLflow documentation pages
• [ ] area/examples: Example code
• [ ] area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• [X] area/models: MLmodel format, model serialization/deserialization, flavors
• [ ] area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• [ ] area/projects: MLproject format, project running backends
• [ ] area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• [ ] area/server-infra: MLflow Tracking server backend
• [ ] area/tracking: Tracking Service, tracking client APIs, autologging

What interface(s) does this bug affect?

• [ ] area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• [ ] 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

What language(s) does this bug affect?

• [ ] language/r: R APIs and clients
• [ ] language/java: Java APIs and clients
• [ ] language/new: Proposals for new client languages

What integration(s) does this bug affect?

• [ ] integrations/azure: Azure and Azure ML integrations
• [ ] integrations/sagemaker: SageMaker integrations
• [ ] integrations/databricks: Databricks integrations

[manijhariya]

I have checked this issue. This issue persist not just for pd.DataFrame but also for other data types also.
If we pass input type of np.ndarray, csc_matrix, csr_matrix, dict of ndarray into infer_signature than it will returns TensorSpec for others it will just return list of Schema for 'ColSpec' which is not compatible to log_model.
Because every field of signature.inputs.inputs expected to be of TensorSpec.

[github-actions[bot]]

@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.