code: def run_matmul_ds(a, b, use_triton_ops=False):
      if use_triton_ops:
          from deepspeed.ops.transformer.inference.triton import (
              matmul_4d as matmul,
          )
  
          return matmul(a, b)
  
      assert use_triton_ops, "Only triton softmax is supported for now"
  
file_path: deepspeed/tests/unit/ops/transformer/inference/test_matmul.py
imports: import pytest
  import torch
  import deepspeed
  from deepspeed.ops.op_builder import InferenceBuilder
label_desc: Use image node_id 3 for calling a global function with example usage: run_matmul_ds(a, b, use_triton_ops) and returns: matmul