docstr: Test fusion support for LSTM (single-layer) -> Linear
docstr_tok: ["Test","fusion","support","for","LSTM","(","single-layer",")","-",">","Linear"]
code: def test_prune_lstm_linear_single_layer(self):
      """
      Test fusion support for LSTM (single-layer) -> Linear
      """
      model = LSTMLinearModel(
          input_dim=8,
          hidden_dim=8,
          output_dim=8,
          num_layers=1,
      )
  
      config = [
          {"tensor_fqn": "lstm.weight_ih_l0"},
          {"tensor_fqn": "lstm.weight_hh_l0"},
      ]
  
      lstm_input = torch.ones((1, 8))
      fx_pruner = BottomHalfLSTMPruner({"sparsity_level": 0.5})
      fx_pruner.prepare(model, config)
      fx_pruner.enable_mask_update = True
      fx_pruner.step()
      model.eval()
  
      out_expected, lstm_out_expected = model(lstm_input)
      pruned_model = fx_pruner.prune()
      pruned_model.eval()
      out_pruned, lstm_out_pruned = pruned_model(lstm_input)
      r, c = lstm_out_expected.size()
  
      # We cannot check that y_expected == y_pruned as usual because
      # zeros vs. missing elements yield different numerical results.
      # Instead that we check that the pruned elements are the first half of the results
      # since we are using a BottomHalfLSTMPruner
      assert torch.isclose(
          lstm_out_expected[:, : c // 2],
          lstm_out_pruned,
          rtol=1e-05,
          atol=1e-07,
      ).all()
      # also check that output of linear is the same shape, this means we've resized
      # linear columns correctly.
      assert out_expected.shape == out_pruned.shape
  
imports: import copy
  import logging
  import random
  import torch
  from torch import nn
  from torch.ao.pruning._experimental.pruner import SaliencyPruner, LSTMSaliencyPruner, BaseStructuredSparsifier, FakeStructuredSparsity, FPGMPruner
  from torch.nn.utils import parametrize
  from torch.testing._internal.common_utils import TestCase, skipIfTorchDynamo
  from torch.testing._internal.common_pruning import SimpleLinear, LinearBias, LinearActivation, LinearActivationFunctional, SimpleConv2d, Conv2dBias, Conv2dActivation, Conv2dPadBias, Conv2dPool, Conv2dPoolFlatten, Conv2dPoolFlattenFunctional, LSTMLinearModel, LSTMLayerNormLinearModel, rows_are_subset
label_desc: Use image node_id 27 for calling the TestBaseStructuredSparsifier obj's underlying member method code with example usage: obj.test_prune_lstm_linear_single_layer() without return types