def METHOD_NAME():
    mx.rnd.seed(seed_state=1234)
    model = DeepTPPPredictionNetwork(
        num_marks=5,
        time_distr_output=WeibullOutput(),
        interval_length=1.0,
        prediction_interval_length=10.0,
    )
    model.initialize()
    past_ia_times = nd.array([[0.1, 0.2, 0.1, 0.12], [0.3, 0.15, 0.1, 0.12]])
    past_marks = nd.array([[1, 2, 0, 2], [0, 0, 1, 2]])
    past_valid_length = nd.array([3, 4])
    past_target = nd.stack(past_ia_times, past_marks, axis=-1)
    pred_target, pred_valid_length = model(past_target, past_valid_length)
    # pred_target must have shape
    # (num_parallel_samples, batch_size, max_sequence_length, 2)
    assert pred_target.ndim == 4
    assert pred_target.shape[0] == model.num_parallel_samples
    assert pred_target.shape[1] == past_ia_times.shape[0]
    assert pred_target.shape[3] == 2  # TPP prediction contains ia_time & mark
    # pred_valid_length must have shape (num_parallel_samples, batch_size)
    assert pred_valid_length.ndim == 2
    assert pred_valid_length.shape[0] == model.num_parallel_samples
    assert pred_valid_length.shape[1] == past_ia_times.shape[0]
    pred_ia_times = pred_target[..., 0].asnumpy()
    pred_marks = pred_target[..., 1].asnumpy()
    assert pred_marks.min() >= 0
    assert pred_marks.max() < model.num_marks
    assert (pred_ia_times >= 0).all()
    # ia_times are set to zero above valid_length (see DeepTPPPredictionNetwork)
    assert (pred_ia_times.sum(-1) < model.prediction_interval_length).all()