code: def test_predict_proba_simple(model, X_masked):
      y_hat = model.predict_proba(X_masked)
  
      assert_array_almost_equal(
          y_hat[0],
          [
              [0.2300, 0.7700],
              [1.0000, 0.0000],
              [0.2011, 0.7989],
              [0.0000, 1.0000],
              [0.1299, 0.8701],
              [1.0000, 0.0000],
              [0.2011, 0.7989],
              [0.0000, 1.0000],
              [0.0000, 1.0000],
              [1.0000, 0.0000],
              [1.0000, 0.0000],
          ],
          4,
      )
  
      assert_array_almost_equal(
          y_hat[1],
          [
              [0.0000, 0.0000, 1.0000],
              [1.0000, 0.0000, 0.0000],
              [0.4469, 0.3453, 0.2078],
              [0.0000, 0.0000, 1.0000],
              [0.0000, 1.0000, 0.0000],
              [0.0000, 1.0000, 0.0000],
              [0.4469, 0.3453, 0.2078],
              [0.2381, 0.4762, 0.2857],
              [0.0000, 1.0000, 0.0000],
              [0.0000, 0.0000, 1.0000],
              [0.5000, 0.2500, 0.2500],
          ],
          4,
      )
  
      assert_array_almost_equal(
          y_hat[2],
          [
              [1.0000, 0.0000],
              [0.3292, 0.6708],
              [0.4916, 0.5084],
              [0.0000, 1.0000],
              [0.4240, 0.5760],
              [1.0000, 0.0000],
              [0.4916, 0.5084],
              [0.0000, 1.0000],
              [1.0000, 0.0000],
              [0.0000, 1.0000],
              [1.0000, 0.0000],
          ],
          4,
      )
  
      assert_array_almost_equal(
          y_hat[3],
          [
              [1.0000, 0.0000],
              [0.5695, 0.4305],
              [1.0000, 0.0000],
              [0.0000, 1.0000],
              [0.5641, 0.4359],
              [0.5312, 0.4688],
              [1.0000, 0.0000],
              [0.0000, 1.0000],
              [0.5312, 0.4688],
              [1.0000, 0.0000],
              [0.5312, 0.4688],
          ],
          4,
      )
  
imports: import numpy
  import torch
  import pytest
  from pomegranate.factor_graph import FactorGraph
  from pomegranate.distributions import Exponential
  from pomegranate.distributions import Categorical
  from pomegranate.distributions import JointCategorical
  from nose.tools import assert_raises
  from numpy.testing import assert_array_almost_equal
label_desc: Use image node_id 21 for calling a global function with example usage: test_predict_proba_simple(model, X_masked) without return types