code: def test_output_layer():
      num_hidden = 32
      vocab_size = 64
      data = pt.ones(2, 10, num_hidden)
      vocab_slice_ids = pt.tensor([4, 7, 23])
  
      b = sockeye.layers.OutputLayer(num_hidden, vocab_size)
      assert b.weight.data.shape == (vocab_size, num_hidden)
  
      output = b(data, None)
      assert output.shape == (2, 10, vocab_size)
      reduced_output = output.index_select(-1, vocab_slice_ids)
  
      output_restricted = b(data, vocab_slice_ids)
      assert output_restricted.shape == (2, 10, len(vocab_slice_ids))
  
      pt.testing.assert_close(
          output_restricted, reduced_output, equal_nan=True
      )
  
file_path: sockeye/test/unit/test_layers.py
imports: from math import pow, sqrt
  import numpy
  import pytest
  import torch
  import sockeye.constants
  import sockeye.layers
  import sockeye.transformer
  from sockeye.knn import KNNConfig, FaissIndexBuilder
label_desc: Use image node_id 4 for calling a global function with example usage: test_output_layer() without return types