Repository: tensorflow/model-analysis
Issue: Unclear EvalConfig nomenclature for model outputs
Root cause: Documentation Deficiency

### System information

-   **Have I written custom code (as opposed to using a stock example script
    provided in TensorFlow Model Analysis)**: Yes
-   **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: macOS Catalina
-   **TensorFlow Model Analysis installed from (source or binary)**: pypi
-   **TensorFlow Model Analysis version (use command below)**: 0.22.1
-   **Python version**: 3.7.5
-   **Jupyter Notebook version**: 1.0.0
-   **TensorFlow version**: 2.3.0

### Describe the problem

I have trained a Keras model with a custom training loop i.e. using gradient tape instead of the `.fit` method. I have specified a custom serving signature as follows:

```python
@tf.function(input_signature=tf.TensorSpec(shape=[None], dtype=tf.string, name="examples"))
def _serving_function(examples):
  mu = model(x)
  sigma = 100.0
  return {"mu": mu, "sigma": sigma}
```

This is because we want the model to output distribution parameters rather than a point estimate. The serving signature as seen by saved_model_cli is:

```
signature_def['serving_default']:
  The given SavedModel SignatureDef contains the following input(s):
    inputs['examples'] tensor_info:
        dtype: DT_STRING
        shape: (-1)
        name: serving_default_examples:0
  The given SavedModel SignatureDef contains the following output(s):
    outputs['mu'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: StatefulPartitionedCall_1:0
    outputs['sigma'] tensor_info:
        dtype: DT_FLOAT
        shape: (-1, 1)
        name: StatefulPartitionedCall_1:1
  Method name is: tensorflow/serving/predict
```

I would like to analyse this model using TFMA, and I have configured it using the following EvalConfig:

```python
eval_config = tfma.EvalConfig(
    model_specs=[tfma.ModelSpec(label_key="my_label_key")],
    metrics_specs=[
        tfma.MetricsSpec(
            metrics=[tfma.MetricConfig(
                class_name="MeanSquaredError",
                config='"dtype": "float32", "name": "mse"'
            )]
        )
    ],
    slicing_specs=[tfma.SlicingSpec()]
)
eval_shared_model = tfma.default_eval_shared_model(saved_model_dir, eval_config=eval_config)
```

However, when I invoke the `run_model_analysis` function, I get the following error:

```python
ValueError: unable to prepare labels and predictions because the labels and/or predictions are dicts with unrecognized keys. If a multi-output keras model (or estimator) was used check that an output_name was provided. If an estimator was used check that common prediction keys were provided (e.g. logistic, probabilities, etc): labels=[1416.], predictions={'mu': array([91.325935], dtype=float32), 'sigma': array([100.], dtype=float32)}, prediction_key= [while running 'ExtractEvaluateAndWriteResults/ExtractAndEvaluate/EvaluateMetricsAndPlots/ComputeMetricsAndPlots()/ComputePerSlice/ComputeUnsampledMetrics/CombinePerSliceKey/WindowIntoDiscarding']
```

This leads me to believe that I might have misconfigured the labels and/or predictions key. However, the [protobuf](https://github.com/tensorflow/model-analysis/blob/v0.23.0/tensorflow_model_analysis/proto/config.proto) is a little vague (to me) as to what those terms refer to. 

What is the difference between labels, predictions, and output? Is my current model a single-output model with different prediction keys, or a multi-output model? How should I be configuring the EvalConfig?

It is a bit confusing. In the case of Estimators, a single output model would output multiple prediction keys (probabilities, classes, etc) but only one of the keys was intended to be used as the prediction (e.g. 'probabilities', 'logistic', etc). If someone customized the name of that key then they would need to update their EvalConfig.model_spec to set the 'prediction_key'. The estimator could also be multi-head (multi-output) in which case there would be multiple dicts each with a set of keys (the 'head' name would correspond to the `EvalConfig.metrics_spec.output_name' and the `EvalConfig.model_spec.prediction_key` would refer to the entry representing the prediction under each head).

For keras there is only a single key per output so only the 'EvalConfig.metrics_spec.output_name' applies for these models.

Long story short, AFAICT you have a multi-output model with two possible outputs 'mu' and 'sigma', so you should set the 'EvalConfig.metrics_spec.output_name' to either (or both) outputs. However, if your setup is such that the prediction is always either 'mu' and 'sigma' and the other key is just extra information, then you should set the 'EvalConfig.model_spec.prediction_key' instead.

Ah, I see. That is rather confusing indeed 😂  What worked for me was the following:

```python
eval_config = tfma.EvalConfig(
    model_specs=[
        tfma.ModelSpec(
            label_key="my_label_key",
            prediction_key="mu" # setting the prediction key
        )
    ],
    metrics_specs=[
        tfma.MetricsSpec(
            output_names=["mu"], # and the output names
            metrics=[tfma.MetricConfig(
                class_name="MeanSquaredError",
                config='"dtype": "float32", "name": "mse"'
            )]
        )
    ],
    slicing_specs=[tfma.SlicingSpec()]
)
eval_shared_model = tfma.default_eval_shared_model(saved_model_dir, eval_config=eval_config)
```

What happens if I want to use a custom metric with a `update_state` method that has a nonstandard function signature? The typical function signature is:

```python
def update_state(self, y_true, y_pred, sample_weight=None):
```

However, my custom metric uses the following signature:

```python
def update_state(self, y_true, y_pred, some, other, arguments, sample_weight=None):
```

I can't seem to figure out (from the TFMA documentation) how to specify a custom metric with a `update_state` that has been overridden. As such, when I run the model analysis, I encounter the following error:

```python
TypeError: in user code:
    TypeError: tf__update_state() missing 3 required positional arguments: 'some', 'other', and 'arguments'
```

Do we only have the option to pass kwargs to initialise the metric via `EvalConfig.MetricConfig.config`?

For keras metrics we currently only support metrics that can be passed to model.compile (i.e. that take the standard arguments: y_true, y_pred, sample_weight). In the future we will support other metrics (i.e. those added via model.add_metric), but those will need to be added as part of your model otherwise we don't know what arguments to pass.

 I understand. Guess I'll close this issue for now since it doesn't seem like we'll be able to get to a resolution. Thanks @mdreves