//  [python]
def test_model_with_multiple_tasks_loss_weights_and_weighted_metrics(
    music_streaming_data: Dataset, run_eagerly: bool
):
    inputs = mm.InputBlockV2(music_streaming_data.schema)
    prediction_tasks = mm.PredictionTasks(music_streaming_data.schema)
    model = mm.Model(inputs, mm.MLPBlock([64]), prediction_tasks)

    loss_weights = {
        "click/binary_classification_task": 1.0,
        "play_percentage/regression_task": 2.0,
        "like/binary_classification_task": 3.0,
    }

    weighted_metrics = {
        "click/binary_classification_task": (
            tf.keras.metrics.Precision(name="precision"),
            tf.keras.metrics.Recall(name="recall"),
            tf.keras.metrics.BinaryAccuracy(name="binary_accuracy"),
            tf.keras.metrics.AUC(name="auc"),
        ),
        "like/binary_classification_task": (
            tf.keras.metrics.Precision(name="weighted_precision"),
            tf.keras.metrics.Recall(name="weighted_recall"),
            tf.keras.metrics.BinaryAccuracy(name="binary_accuracy"),
            tf.keras.metrics.AUC(name="weighted_auc"),
        ),
        "play_percentage/regression_task": (
            tf.keras.metrics.RootMeanSquaredError(name="root_mean_squared_error"),
        ),
    }

    model.compile(
        optimizer="adam",
        run_eagerly=run_eagerly,
        loss_weights=loss_weights,
        weighted_metrics=weighted_metrics,
    )

    batch = mm.sample_batch(music_streaming_data, batch_size=50, prepare_features=False)

    metrics = model.test_step(batch)

    assert metrics["loss"] >= 0
    assert set(metrics.keys()) == set(
        [
            "loss",
            "click/binary_classification_task_loss",
            "like/binary_classification_task_loss",
            "play_percentage/regression_task_loss",
            "click/binary_classification_task_precision",
            "click/binary_classification_task_recall",
            "click/binary_classification_task_binary_accuracy",
            "click/binary_classification_task_auc",
            "click/binary_classification_task_weighted_precision",
            "click/binary_classification_task_weighted_recall",
            "click/binary_classification_task_weighted_binary_accuracy",
            "click/binary_classification_task_weighted_auc",
            "like/binary_classification_task_precision",
            "like/binary_classification_task_recall",
            "like/binary_classification_task_binary_accuracy",
            "like/binary_classification_task_auc",
            "like/binary_classification_task_weighted_precision",
            "like/binary_classification_task_weighted_recall",
            "like/binary_classification_task_weighted_binary_accuracy",
            "like/binary_classification_task_weighted_auc",
            "play_percentage/regression_task_root_mean_squared_error",
            "play_percentage/regression_task_weighted_root_mean_squared_error",
            "regularization_loss",
            "loss_batch",
        ]
    )

    metrics2 = model.test_step(batch)
    assert metrics["loss"] == metrics2["loss"]
    for m in metrics:
        assert metrics[m] == metrics2[m]

    # Disabling losses weights
    model.compile(loss_weights=None)
    metrics_non_weighted = model.test_step(batch)

    assert metrics_non_weighted["loss"] != metrics["loss"]
    assert (
        metrics_non_weighted["click/binary_classification_task_loss"]
        == metrics["click/binary_classification_task_loss"]
    )
    assert (
        metrics_non_weighted["like/binary_classification_task_loss"]
        == metrics["like/binary_classification_task_loss"]
    )