code: def test_one_hot_custom_layout_axis_name():
      np.random.seed(42)
      layout = get_initial_layout(3)
      for axis_name in "Xx01":
          yield partial(
              check_one_hot_operator,
              axis=-1,
              initial_layout=layout,
              axis_name=axis_name,
          ), random_3d_tensors_batch
  
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 15 for calling a global function with example usage: test_one_hot_custom_layout_axis_name() without return types