code: def test_softmax(batch, sequence, channels, dtype, use_triton_ops):
      if not deepspeed.HAS_TRITON and use_triton_ops:
          pytest.skip("triton has to be installed for the test")
      input_ds = torch.randn(
          (batch, sequence, channels), dtype=dtype, device="cuda"
      )
      input_ref = input_ds.clone().detach()
  
      ds_out = run_softmax_ds(input_ds, use_triton_ops)
      ref_out = run_softmax_reference(input_ref)
      assert allclose(ds_out, ref_out)
  
imports: import pytest
  import torch
  import deepspeed
  from deepspeed.ops.op_builder import InferenceBuilder
label_desc: Use image node_id 4 for calling a global function with example usage: test_softmax(batch, sequence, channels, dtype, use_triton_ops) without return types