code: def test_hidden_layer_log_prob(
      non_linearity, include_hidden_bias, B=2, D=3, H=2
  ):
      X = torch.randn(B, D)
      A_mean = torch.rand(D, H)
      A_scale = 0.3 * torch.exp(0.3 * torch.rand(D, H))
      dist = HiddenLayer(
          X=X,
          A_mean=A_mean,
          A_scale=A_scale,
          non_linearity=non_linearity,
          include_hidden_bias=include_hidden_bias,
      )
  
      A_dist = Normal(A_mean, A_scale)
      A_prior = Normal(torch.zeros(D, H), torch.ones(D, H))
      kl = torch.distributions.kl.kl_divergence(A_dist, A_prior).sum()
      assert_equal(kl, dist.KL, prec=0.01)
  
file_path: pyro/tests/contrib/bnn/test_hidden_layer.py
imports: import pytest
  import torch
  import torch.nn.functional
  from torch.distributions import Normal
  from pyro.contrib.bnn import HiddenLayer
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_hidden_layer_log_prob(non_linearity, include_hidden_bias, B, D, H) without return types