code: def test_prediction_network_output():
      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()
  
file_path: gluonts/test/mx/model/tpp/test_deeptpp.py
imports: import mxnet
  import numpy
  import pytest
  from mxnet import nd
  from gluonts.mx.model.tpp.deeptpp._network import DeepTPPPredictionNetwork, DeepTPPTrainingNetwork
  from gluonts.mx.model.tpp.distribution import WeibullOutput
label_desc: Use image node_id 5 for calling a global function with example usage: test_prediction_network_output() without return types