code: def test_prediction_result_with_proba(self):
      proba = np.array(
          [
              [0.1, 0.2, 0.6, 0.1],
              [0.25, 0.25, 0.25, 0.25],
              [0.3, 0.3, 0.2, 0.2],
              [0.3, 0.2, 0.5, 0.1],
          ]
      )
      result, proba_result = prediction_result(
          proba, False, proba.shape[1], return_proba=True
      )
      expected = np.array([2, 0, 0, 2])
      assert_array_equal(expected, result)
      assert_array_equal(proba, proba_result)
  
file_path: small-text/tests/unit/small_text/utils/test_classification.py
imports: import unittest
  import numpy
  from numpy.testing import assert_array_equal
  from scipy.sparse import csr_matrix
  from sklearn.preprocessing import MultiLabelBinarizer
  from small_text.utils.classification import empty_result, prediction_result
  from tests.utils.testing import assert_csr_matrix_equal
label_desc: Use image node_id 2 for calling the ClassificationUtilsTest obj's underlying member method code with example usage: obj.test_prediction_result_with_proba() without return types