code: def test_multi_headed_mlp_module(
      input_dim, output_dim, hidden_sizes, output_w_init_vals, n_heads
  ):
      """Test Multi-headed MLPModule.
  
      Args:
          input_dim (int): Input dimension.
          output_dim (int): Ouput dimension.
          hidden_sizes (list[int]): Size of hidden layers.
          output_w_init_vals (list[int]): Init values for output weights.
          n_heads (int): Number of output layers.
  
      """
      module = MultiHeadedMLPModule(
          n_heads=n_heads,
          input_dim=input_dim,
          output_dims=output_dim,
          hidden_sizes=hidden_sizes,
          hidden_nonlinearity=None,
          hidden_w_init=nn.init.ones_,
          output_nonlinearities=None,
          output_w_inits=list(
              map(_helper_make_inits, output_w_init_vals)
          ),
      )
  
      input_value = torch.ones(input_dim)
      outputs = module(input_value)
  
      if len(output_w_init_vals) == 1:
          output_w_init_vals = list(output_w_init_vals) * n_heads
      if len(output_dim) == 1:
          output_dim = list(output_dim) * n_heads
      for i, output in enumerate(outputs):
          expected = input_dim * torch.Tensor(hidden_sizes).prod()
          expected *= output_w_init_vals[i]
          assert torch.equal(
              output,
              torch.full((output_dim[i],), expected, dtype=torch.float),
          )
  
imports: import pytest
  import torch
  from torch import nn
  from garage.torch.modules import MultiHeadedMLPModule
label_desc: Use image node_id 2 for calling a global function with example usage: test_multi_headed_mlp_module(input_dim, output_dim, hidden_sizes, output_w_init_vals, n_heads) without return types