code: def check_one_hot_operator(
      source,
      device="cpu",
      axis=-1,
      expected_output_dim=None,
      axis_name=None,
      initial_layout=None,
  ):
      pipeline = OneHotPipeline(
          num_classes=num_classes,
          source=source,
          axis=axis,
          layout=initial_layout,
          axis_name=axis_name,
          device=device,
      )
      pipeline.build()
      (outputs, input_batch) = pipeline.run()
      if device == "gpu":
          input_batch = input_batch.as_cpu()
      input_batch = list(map(np.array, input_batch))
      expected_output_dim = (
          expected_output_dim or len(input_batch[0].shape) + 1
      )
      reference = (
          one_hot_3_axes(input_batch, axis)
          if expected_output_dim == 4
          else one_hot(input_batch)
      )
      expected_layout = modify_layout(
          initial_layout, expected_output_dim, axis, axis_name
      )
      check_batch(
          outputs,
          reference,
          batch_size,
          max_allowed_error=0,
          expected_layout=expected_layout,
      )
  
imports: from functools import partial
  from nvidia.dali.pipeline import Pipeline
  import nvidia.dali.ops
  import nvidia.dali.types
  from test_utils import check_batch
  from nose_utils import raises
  import numpy
label_desc: Use image node_id 9 for calling a global function with example usage: check_one_hot_operator(source, device, axis, expected_output_dim, axis_name, initial_layout) without return types