code: def get_data(
      shape: Tuple[Tuple[int, int], Tuple[int, int]],
      dtype: torch.dtype = torch.float16,
      device: str = "cuda",
  ) -> Tuple[torch.Tensor, nn.Parameter, torch.Tensor]:
      """Utility function for getting some tensors for testing and benchmarking."""
      (tokens, d1), (d2, vocabs) = shape
      assert d1 == d2
      input = torch.rand(
          tokens, d1, device=device, dtype=dtype
      ).requires_grad_(True)
      # Before pytorch 1.9, nn.Linear does not support device and dtype init option. So we use to()
      # and an if condition.
      layer = nn.Linear(d2, vocabs, bias=False).to(device)
      assert dtype in [torch.float16, torch.float32]
      if dtype == torch.float16:
          layer = layer.half()
      weight = layer.weight
      target = (torch.rand(tokens, device=device) * vocabs).long()
      return input, weight, target
  
imports: from typing import Any, Optional, Tuple
  import torch
  from torch import nn
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 3 for calling a global function with example usage: get_data(shape, dtype, device) and returns: input,  weight,  target