code: def test_get_softmax_arg_model(
      data: np.ndarray,
      params: dict,
      expected: np.ndarray,
      load_mnist_model_tf,
  ):
      model = TensorFlowModel(model=load_mnist_model_tf, softmax=True)
      new_model = TensorFlowModel(
          model=model.get_softmax_arg_model(), **params
      )
  
      out = new_model.predict(x=data)
  
      assert np.allclose(out, expected), "Test failed."
  
file_path: Quantus/tests/helpers/test_tf_model.py
imports: from functools import reduce
  from operator import and_
  from typing import Union
  import tensorflow
  import numpy
  import pytest
  from pytest_lazyfixture import lazy_fixture
  from scipy.special import softmax
  from quantus.helpers.model.tf_model import TensorFlowModel
label_desc: Use image node_id 1 for calling a global function with example usage: test_get_softmax_arg_model(data, params, expected, load_mnist_model_tf) without return types