code: def bayesian_linear_model(
      design,
      w_means={},
      w_sqrtlambdas={},
      re_group_sizes={},
      re_alphas={},
      re_betas={},
      obs_sd=None,
      alpha_0=None,
      beta_0=None,
      response="normal",
      response_label="y",
      k=None,
  ):
      """
      A pyro model for Bayesian linear regression.
  
      If :param:`response` is `"normal"` this corresponds to a linear regression
      model
  
          :math:`Y = Xw + \\epsilon`
  
      with `\\epsilon`` i.i.d. zero-mean Gaussian. The observation standard deviation
      (:param:`obs_sd`) may be known or unknown. If unknown, it is assumed to follow an
      inverse Gamma distribution with parameters :param:`alpha_0` and :param:`beta_0`.
  
      If the response type is `"bernoulli"` we instead have :math:`Y \\sim Bernoulli(p)`
      with
  
          :math:`logit(p) = Xw`
  
      Given parameter groups in :param:`w_means` and :param:`w_sqrtlambda`, the fixed effects
      regression coefficient is taken to be Gaussian with mean `w_mean` and standard deviation
      given by
  
          :math:`\\sigma / \\sqrt{\\lambda}`
  
      corresponding to the normal inverse Gamma family.
  
      The random effects coefficient is constructed as follows. For each random effect
      group, standard deviations for that group are sampled from a normal inverse Gamma
      distribution. For each group, a random effect coefficient is then sampled from a zero
      mean Gaussian with those standard deviations.
  
      :param torch.Tensor design: a tensor with last two dimensions `n` and `p`
              corresponding to observations and features respectively.
      :param OrderedDict w_means: map from variable names to tensors of fixed effect means.
      :param OrderedDict w_sqrtlambdas: map from variable names to tensors of square root
          :math:`\\lambda` values for fixed effects.
      :param OrderedDict re_group_sizes: map from variable names to int representing the
          group size
      :param OrderedDict re_alphas: map from variable names to `torch.Tensor`, the tensor
          consists of Gamma dist :math:`\\alpha` values
      :param OrderedDict re_betas: map from variable names to `torch.Tensor`, the tensor
          consists of Gamma dist :math:`\\beta` values
      :param torch.Tensor obs_sd: the observation standard deviation (if assumed known).
          This is still relevant in the case of Bernoulli observations when coefficeints
          are sampled using `w_sqrtlambdas`.
      :param torch.Tensor alpha_0: Gamma :math:`\\alpha` parameter for unknown observation
          covariance.
      :param torch.Tensor beta_0: Gamma :math:`\\beta` parameter for unknown observation
          covariance.
      :param str response: Emission distribution. May be `"normal"` or `"bernoulli"`.
      :param str response_label: Variable label for response.
      :param torch.Tensor k: Only used for a sigmoid response. The slope of the sigmoid
          transformation.
      """
      # design is size batch x n x p
      # tau is size batch
      batch_shape = design.shape[:-2]
      with ExitStack() as stack:
          for plate in iter_plates_to_shape(batch_shape):
              stack.enter_context(plate)
  
          if obs_sd is None:
              # First, sample tau (observation precision)
              tau_prior = dist.Gamma(
                  alpha_0.unsqueeze(-1), beta_0.unsqueeze(-1)
              ).to_event(1)
              tau = pyro.sample("tau", tau_prior)
              obs_sd = 1.0 / torch.sqrt(tau)
  
          elif alpha_0 is not None or beta_0 is not None:
              warnings.warn(
                  "Values of `alpha_0` and `beta_0` unused becased"
                  "`obs_sd` was specified already."
              )
  
          obs_sd = obs_sd.expand(batch_shape + (1,))
  
          # Build the regression coefficient
          w = []
          # Allow different names for different coefficient groups
          # Process fixed effects
          for name, w_sqrtlambda in w_sqrtlambdas.items():
              w_mean = w_means[name]
              # Place a normal prior on the regression coefficient
              w_prior = dist.Normal(
                  w_mean, obs_sd / w_sqrtlambda
              ).to_event(1)
              w.append(pyro.sample(name, w_prior))
          # Process random effects
          for name, group_size in re_group_sizes.items():
              # Sample `G` once for this group
              alpha, beta = re_alphas[name], re_betas[name]
              G_prior = dist.Gamma(alpha, beta).to_event(1)
              G = 1.0 / torch.sqrt(pyro.sample("G_" + name, G_prior))
              # Repeat `G` for each group
              repeat_shape = tuple(1 for _ in batch_shape) + (
                  group_size,
              )
              u_prior = dist.Normal(
                  torch.tensor(0.0), G.repeat(repeat_shape)
              ).to_event(1)
              w.append(pyro.sample(name, u_prior))
          # Regression coefficient `w` is batch x p
          w = broadcast_cat(w)
  
          # Run the regressor forward conditioned on inputs
          prediction_mean = rmv(design, w)
          if response == "normal":
              # y is an n-vector: hence use .to_event(1)
              return pyro.sample(
                  response_label,
                  dist.Normal(prediction_mean, obs_sd).to_event(1),
              )
          elif response == "bernoulli":
              return pyro.sample(
                  response_label,
                  dist.Bernoulli(logits=prediction_mean).to_event(1),
              )
          elif response == "sigmoid":
              base_dist = dist.Normal(prediction_mean, obs_sd).to_event(
                  1
              )
              # You can add loc via the linear model itself
              k = k.expand(prediction_mean.shape)
              transforms = [
                  AffineTransform(loc=torch.tensor(0.0), scale=k),
                  SigmoidTransform(),
              ]
              response_dist = dist.TransformedDistribution(
                  base_dist, transforms
              )
              return pyro.sample(response_label, response_dist)
          else:
              raise ValueError(
                  "Unknown response distribution: '{}'".format(response)
              )
  
imports: import warnings
  from collections import OrderedDict
  from contextlib import ExitStack
  from functools import partial
  import torch
  from torch.distributions import constraints
  from torch.distributions.transforms import AffineTransform, SigmoidTransform
  from torch.nn.functional import softplus
  import pyro
  import pyro.distributions
  from pyro.contrib.util import iter_plates_to_shape, rmv
label_desc: Use image node_id 11 for calling a global function with example usage: bayesian_linear_model(design, w_means, w_sqrtlambdas, re_group_sizes, re_alphas, re_betas, obs_sd, alpha_0, beta_0, response, response_label, k) and returns: pyro, pyro, pyro