code: def make_prior(k):
      assert 0 < k <= 1
      u = 1 / (1 + k + k**2 + k**3)
      p0 = 1 - u
      p1 = 1 - (k * u) / p0
      p2 = 1 - (k**2 * u) / (p0 * p1)
      trial_probs = [p0, p1, p2]
      # dist = [1 - p0, p0 * (1 - p1), p0 * p1 * (1 - p2), p0 * p1 * p2]
      # print(dist)
      return lambda t: trial_probs[t]
  
imports: import argparse
  import math
  import os
  import time
  from functools import partial
  import numpy
  import torch
  import visdom
  from air import AIR, latents_to_tensor
  from viz import draw_many, tensor_to_objs
  import pyro
  import pyro.contrib.examples.multi_mnist
  import pyro.optim
  import pyro.poutine
  from pyro.contrib.examples.util import get_data_directory
  from pyro.infer import SVI, JitTraceGraph_ELBO, TraceGraph_ELBO
label_desc: Use image node_id 2 for calling a global function with example usage: make_prior(k) and returns: function