code: def _init_matmul_argument(
      rows: int, cols: int, quantization: QuantizedInputType
  ) -> torch.Tensor | triton.TensorWrapper:
      """Initialize an input for matrix multiplication."""
      if quantization == QuantizedInputType.INT8:
          return torch.randint(
              0, 128, (rows, cols), device="cuda", dtype=torch.int8
          )
      elif quantization == QuantizedInputType.FLOAT8:
          return to_triton_f8(
              torch.randn(
                  rows, cols, device="cuda", dtype=torch.bfloat16
              )
          )
      elif quantization == QuantizedInputType.FLOAT16:
          return torch.randn(
              rows, cols, device="cuda", dtype=torch.float16
          )
      else:
          return torch.randn(
              rows, cols, device="cuda", dtype=torch.bfloat16
          )
  
imports: import enum
  import itertools
  import logging
  import math
  import typing
  import torch
  import tqdm
  import triton
  import triton.language
label_desc: Use image node_id 11 for calling a global function with example usage: _init_matmul_argument(rows, cols, quantization) and returns: torch, to_triton_f8, torch, torch