code: def test_rnn_input(
      prediction_length: int, context_length: int, lags_seq: List[int]
  ):
      num_samples = 10
      history_length = max(lags_seq) + context_length - 1
  
      model = DeepARModel(
          freq="D",
          prediction_length=prediction_length,
          context_length=context_length,
          lags_seq=lags_seq,
          scaling=False,
          num_parallel_samples=num_samples,
      )
  
      # Construct test batch of size 1, such that:
      # - target values are increasing integers
      # - there is one dynamic feature that is equal to the target
      #
      # Since scaling=False, this way we can compare lagged target
      # values and dynamic feature and verify that the expected values
      # are given as input to the RNN, at the expected time index.
  
      batch = {
          "feat_static_cat": torch.tensor([[0]], dtype=torch.int64),
          "feat_static_real": torch.tensor(
              [[0.0]], dtype=torch.float32
          ),
          "past_time_feat": torch.arange(
              history_length, dtype=torch.float32
          ).view(1, history_length, 1),
          "past_target": torch.arange(
              history_length, dtype=torch.float32
          ).view(1, history_length),
          "past_observed_values": torch.arange(
              history_length, dtype=torch.float32
          ).view(1, history_length),
      }
  
      # test with no future_target (only one step prediction)
  
      batch["future_time_feat"] = torch.arange(
          history_length,
          history_length + 1,
          dtype=torch.float32,
      ).view(1, 1, 1)
  
      rnn_input, scale, _ = model.prepare_rnn_input(**batch)
  
      assert (scale == 1.0).all()
  
      ref = torch.arange(
          history_length - context_length + 1,
          history_length + 1,
          dtype=torch.float32,
      )
  
      for idx, lag in enumerate(lags_seq):
          assert torch.equal(ref - lag, rnn_input[0, :, idx])
  
      # test with all future data
  
      batch["future_time_feat"] = torch.arange(
          history_length,
          history_length + prediction_length,
          dtype=torch.float32,
      ).view(1, prediction_length, 1)
  
      batch["future_target"] = torch.arange(
          history_length,
          history_length + prediction_length,
          dtype=torch.float32,
      ).view(1, prediction_length)
  
      rnn_input, scale, _ = model.prepare_rnn_input(**batch)
  
      assert (scale == 1.0).all()
  
      ref = torch.arange(
          history_length - context_length + 1,
          history_length + prediction_length,
          dtype=torch.float32,
      )
  
      for idx, lag in enumerate(lags_seq):
          assert torch.equal(ref - lag, rnn_input[0, :, idx])
  
imports: from typing import List, Optional
  import pytest
  import torch
  from gluonts.torch.model.deepar import DeepARLightningModule, DeepARModel
label_desc: Use image node_id 2 for calling a global function with example usage: test_rnn_input(prediction_length, context_length, lags_seq) without return types