code: def test_single_input(step_class, teardown):
      x = Input()
      y = step_class()(x)
      model = Model(x, y)
  
      x_data = np.array([[1, 2], [3, 4]])
      if step_class is Stack:
          assert_array_equal(
              x_data.reshape((2, 2, 1)), model.predict(x_data)
          )
      else:
          assert_array_equal(x_data, model.predict(x_data))
  
file_path: baikal/tests/steps/test_merge.py
imports: import numpy
  import pytest
  from numpy.testing import assert_array_equal
  from baikal import Model, Input
  from baikal._core.utils import listify, safezip2
  from baikal.steps import ColumnStack, Concatenate, Stack, Split
  from tests.helpers.fixtures import teardown
label_desc: Use image node_id 5 for calling a global function with example usage: test_single_input(step_class, teardown) without return types