code: def test_cross_entropy_layer_backward():
      batch_size, seq_len = kt.bs_sl()
      vocab_size = random.randint(30413, 40519)
      print(
          f"(batch_size, seq_len, vocab_size): ({batch_size}, {seq_len}, {vocab_size})"
      )
  
      base_inputs = kt.rand(
          (batch_size, seq_len, vocab_size)
      ).requires_grad_()
      cus_inputs = base_inputs.clone().detach().requires_grad_()
      targets = kt.randint(
          ce_config_fp16.padding_idx - 1,
          vocab_size,
          (batch_size, seq_len),
      )
      targets_32 = targets.to(torch.int32)
  
      if kt.dtype == torch.float:
          custom_layer = custom_cross_entropy_layer_fp32
      else:
          custom_layer = custom_cross_entropy_layer_fp16
      cus_res = custom_layer(cus_inputs, targets_32)[0].to(kt.dtype)
      x = torch.nn.functional.log_softmax(
          base_inputs, dim=-1, dtype=torch.float32
      )
      base_res, _ = label_smoothed_nll_loss(
          x,
          targets,
          ce_config_fp16.epsilon,
          ignore_index=ce_config_fp16.padding_idx,
      )
      base_res = base_res.to(kt.dtype)
  
      def custom():
          if cus_inputs.grad is not None:
              cus_inputs.grad.zero_()
          cus_res.backward(retain_graph=True)
          return [
              cus_inputs.grad.contiguous().detach(),
          ]
  
      def baseline():
          if base_inputs.grad is not None:
              base_inputs.grad.zero_()
          base_res.backward(retain_graph=True)
          return [
              base_inputs.grad.contiguous().detach(),
          ]
  
      return custom, baseline
  
imports: import random
  import math
  from copy import deepcopy
  from dataclasses import dataclass
  import torch
  import torch.nn
  from tests.util import TestDecorator, get_fairseq_enc_params, get_fairseq_dec_params, max_batch_tokens, max_seq_len, split_custom_layer_grad, copy_grad_from_paras, copy_cmax_grad_from_paras
  from tests import fairseq_layers
  from lightseq.training.ops.pytorch.transformer_encoder_layer import LSTransformerEncoderLayer
  from lightseq.training.ops.pytorch.transformer_embedding_layer import LSTransformerEmbeddingLayer
  from lightseq.training.ops.pytorch.cross_entropy_layer import LSCrossEntropyLayer
  from lightseq.training.ops.pytorch.quant_linear_layer import LSQuantLinearLayer
  from lightseq.training.cli.fs_modules.ls_fs_transformer_decoder_layer import LSFSTransformerDecoderLayer
  from lightseq.training.ops.pytorch.quantization import enable_quant, disable_quant, qat_mode, ptq_mode
label_desc: Use image node_id 30 for calling a global function with example usage: test_cross_entropy_layer_backward() and returns: custom,  baseline, list, list