code: def test_tabnet(
      entmax_mode: Optional[str],
      input_size: int,
      output_size: int,
      size: int,
      virtual_batch_size: Optional[int],
      batch_size: int,
  ) -> None:
      # setup synthetic tensor
      torch.manual_seed(RANDOM_SEED)
      input_tensor = torch.randn(
          [batch_size, input_size], dtype=torch.float32
      )
  
      tabnet = TabNet(
          input_size,
          size,
          output_size,
          num_steps=3,
          num_total_blocks=4,
          num_shared_blocks=2,
          entmax_mode=entmax_mode,
      )
  
      output = tabnet(input_tensor)
  
      # check for expected shape and properties
      assert isinstance(output, tuple)
      assert output[0].shape == (batch_size, output_size)
  
      assert tabnet.input_shape[-1] == input_size
      assert tabnet.output_shape[-1] == output_size
      assert tabnet.input_dtype == torch.float32
  
      # check for parameter updates
      target = torch.randn([batch_size, 1])
      fpc, tpc, upc, not_updated = check_module_parameters_updated(
          tabnet, (input_tensor,), target
      )
  
      if batch_size == 1:
          # for single record batches, batchnorm layer is bypassed, only a subset of parameters are updated
          assert upc == 17, (
              f"Updated parameter count not expected value. Parameters not updated: {not_updated}"
              f"\nModule structure:\n{tabnet}"
          )
      else:
          # update count should equal trainable number of parameters
          assert tpc == upc, (
              f"All parameter not updated. Parameters not updated: {not_updated}"
              f"\nModule structure:\n{tabnet}"
          )
  
file_path: ludwig/tests/ludwig/modules/test_tabnet_modules.py
imports: from typing import Optional
  import pytest
  import torch
  from ludwig.modules.tabnet_modules import AttentiveTransformer, FeatureBlock, FeatureTransformer, TabNet
  from ludwig.utils.entmax import sparsemax
  from tests.integration_tests.parameter_update_utils import check_module_parameters_updated
label_desc: Use image node_id 5 for calling a global function with example usage: test_tabnet(entmax_mode, input_size, output_size, size, virtual_batch_size, batch_size) without return types