//  [python]
def test_mmoe_block_task_specific_sample_weight_and_weighted_metrics(
    music_streaming_data: Dataset, run_eagerly: bool
):
    class CustomSampleWeight(Block):
        def call_outputs(
            self,
            outputs: PredictionOutput,
            features: Dict[str, tf.Tensor] = None,
            targets: tf.Tensor = None,
            training=True,
            testing=False,
            **kwargs,
        ) -> PredictionOutput:
            # Computes loss for the like loss only for clicked items
            outputs = outputs.copy_with_updates(
                sample_weight=tf.expand_dims(targets["click"], -1),
            )
            return outputs

    inputs = mm.InputBlockV2(music_streaming_data.schema)
    prediction_tasks = mm.PredictionTasks(
        music_streaming_data.schema, task_pre_dict={"like": CustomSampleWeight()}
    )
    mmoe = mm.MMOEBlock(prediction_tasks, expert_block=mm.MLPBlock([64]), num_experts=4)
    model = mm.Model(inputs, mmoe, prediction_tasks)

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

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

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

    metrics = model.train_step(
        mm.sample_batch(music_streaming_data, batch_size=50, prepare_features=False)
    )

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