code: def test_predict_proba_one_edge():
      m = Categorical([[1.0 / 3, 1.0 / 3, 1.0 / 3]])
      f = Categorical([[0.23, 0.17, 0.60]])
  
      X = torch.tensor([[0], [0], [1], [2]])
      mask = torch.tensor([[False], [True], [True], [True]])
      X_masked = torch.masked.MaskedTensor(X, mask=mask)
  
      model = FactorGraph([f], [m], [(m, f)])
      y_hat = model.predict_proba(X_masked)
  
      assert_array_almost_equal(
          y_hat[0],
          [
              [0.23, 0.17, 0.60],
              [1.00, 0.00, 0.00],
              [0.00, 1.00, 0.00],
              [0.00, 0.00, 1.00],
          ],
      )
  
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 17 for calling a global function with example usage: test_predict_proba_one_edge() without return types