code: def test_tft_modules(
      d_past_feat_dynamic_real: List[int],
      c_past_feat_dynamic_cat: List[int],
      d_feat_dynamic_real: List[int],
      c_feat_dynamic_cat: List[int],
      d_feat_static_real: List[int],
      c_feat_static_cat: List[int],
      quantiles: List[float],
  ):
      batch_size = 4
      prediction_length = 6
      context_length = 12
  
      lightning_module = TemporalFusionTransformerLightningModule(
          {
              "context_length": context_length,
              "prediction_length": prediction_length,
              "d_past_feat_dynamic_real": d_past_feat_dynamic_real,
              "c_past_feat_dynamic_cat": c_past_feat_dynamic_cat,
              "d_feat_dynamic_real": d_feat_dynamic_real,
              "c_feat_dynamic_cat": c_feat_dynamic_cat,
              "d_feat_static_real": d_feat_static_real,
              "c_feat_static_cat": c_feat_static_cat,
              "distr_output": QuantileOutput(quantiles),
          }
      )
      model = lightning_module.model
  
      feat_static_cat = torch.zeros(
          batch_size, len(c_feat_static_cat), dtype=torch.long
      )
      feat_static_real = torch.ones(batch_size, sum(d_feat_static_real))
      feat_dynamic_cat = torch.zeros(
          batch_size,
          prediction_length + context_length,
          len(c_feat_dynamic_cat),
          dtype=torch.long,
      )
      feat_dynamic_real = torch.ones(
          batch_size,
          prediction_length + context_length,
          sum(d_feat_dynamic_real),
      )
      past_feat_dynamic_cat = torch.zeros(
          batch_size,
          context_length,
          len(c_past_feat_dynamic_cat),
          dtype=torch.long,
      )
      past_feat_dynamic_real = torch.ones(
          batch_size, context_length, sum(d_past_feat_dynamic_real)
      )
      past_target = torch.ones(batch_size, context_length)
      past_observed_values = torch.ones(batch_size, context_length)
      future_target = torch.ones(batch_size, prediction_length)
      future_observed_values = torch.ones(batch_size, prediction_length)
  
      (
          past_covariates,
          future_covariates,
          static_covariates,
          loc,
          scale,
      ) = model._preprocess(
          past_target=past_target,
          past_observed_values=past_observed_values,
          feat_static_real=feat_static_real,
          feat_static_cat=feat_static_cat,
          feat_dynamic_real=feat_dynamic_real,
          feat_dynamic_cat=feat_dynamic_cat,
          past_feat_dynamic_real=past_feat_dynamic_real,
          past_feat_dynamic_cat=past_feat_dynamic_cat,
      )
      for x in past_covariates:
          assert x.shape == (batch_size, context_length, model.d_var)
      for x in future_covariates:
          assert x.shape == (batch_size, prediction_length, model.d_var)
      for x in static_covariates:
          assert x.shape == (batch_size, model.d_var)
      assert loc.shape == scale.shape == (batch_size, 1)
  
      (output,), loc, scale = model(
          past_target=past_target,
          past_observed_values=past_observed_values,
          feat_static_real=feat_static_real,
          feat_static_cat=feat_static_cat,
          feat_dynamic_real=feat_dynamic_real,
          feat_dynamic_cat=feat_dynamic_cat,
          past_feat_dynamic_real=past_feat_dynamic_real,
          past_feat_dynamic_cat=past_feat_dynamic_cat,
      )
  
      assert output.shape == (
          batch_size,
          prediction_length,
          len(quantiles),
      )
      assert loc.shape == scale.shape == (batch_size, 1)
  
      batch = dict(
          past_target=past_target,
          past_observed_values=past_observed_values,
          past_feat_dynamic_real=past_feat_dynamic_real,
          past_feat_dynamic_cat=past_feat_dynamic_cat,
          feat_dynamic_real=feat_dynamic_real,
          feat_dynamic_cat=feat_dynamic_cat,
          feat_static_real=feat_static_real,
          feat_static_cat=feat_static_cat,
          future_target=future_target,
          future_observed_values=future_observed_values,
      )
  
      assert (
          lightning_module.training_step(batch, batch_idx=0).shape == ()
      )
      assert (
          lightning_module.validation_step(batch, batch_idx=0).shape
          == ()
      )
  
imports: from typing import List
  import torch
  import pytest
  from gluonts.torch.distributions import QuantileOutput
  from gluonts.torch.model.tft import TemporalFusionTransformerLightningModule
label_desc: Use image node_id 1 for calling a global function with example usage: test_tft_modules(d_past_feat_dynamic_real, c_past_feat_dynamic_cat, d_feat_dynamic_real, c_feat_dynamic_cat, d_feat_static_real, c_feat_static_cat, quantiles) without return types