code: def test_predict_proba_one_edge(distributions):
      d1, d2, d3, d4, d12, d22, d32, d13 = distributions
  
      model = BayesianNetwork([d1, d22], [(d1, d22)])
  
      X = torch.tensor([[0, 0], [0, 1], [1, 1], [1, 2]])
      mask = torch.tensor(
          [[False, False], [True, False], [True, True], [False, True]]
      )
      X_masked = torch.masked.MaskedTensor(X, mask=mask)
  
      y_hat = model.predict_proba(X_masked)
  
      assert_array_almost_equal(
          y_hat[0],
          [
              [0.8889, 0.1111],
              [1.0000, 0.0000],
              [0.0000, 1.0000],
              [0.9412, 0.0588],
          ],
          4,
      )
  
      assert_array_almost_equal(
          y_hat[1],
          [
              [0.6778, 0.1333, 0.1889],
              [0.7000, 0.1000, 0.2000],
              [0.0000, 1.0000, 0.0000],
              [0.0000, 0.0000, 1.0000],
          ],
          4,
      )
  
imports: import numpy
  import torch
  import pytest
  from pomegranate.bayesian_network import BayesianNetwork
  from pomegranate.distributions import Exponential
  from pomegranate.distributions import Categorical
  from pomegranate.distributions import ConditionalCategorical
  from nose.tools import assert_raises
  from numpy.testing import assert_array_equal
  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(distributions) without return types