code: def _compute_mmd(X, Z, kernel):
      mmd = (
          torch.mean(kernel(X))
          + torch.mean(kernel(Z))
          - torch.mean(kernel(X, Z)) * 2
      )
      return mmd
  
imports: from collections import defaultdict
  import torch
  import pyro
  import pyro.ops.jit
  from pyro import poutine
  from pyro.infer.elbo import ELBO
  from pyro.infer.enum import get_importance_trace
  from pyro.infer.util import is_validation_enabled, torch_item
  from pyro.util import check_if_enumerated, warn_if_nan
label_desc: Use image node_id 1 for calling a global function with example usage: _compute_mmd(X, Z, kernel) and returns: mmd