code: def tree_sample(
      theta,
      logp,
      r0,
      grad,
      extra,
      epsilon,
      invmassmatrix,
      f,
      joint,
      maxheight=np.inf,
  ):
      """Build NUTS-like tree of sampling path from theta towards p with stepsize epsilon."""
      # initialize the tree
      thetaminus = theta
      thetaplus = theta
      rminus = r0[:]
      rplus = r0[:]
      gradminus = grad[:]
      gradplus = grad[:]
      alpha = 1
      beta = 1
      nalpha = 1
  
      j = 0  # initial heigth j = 0
      s = True  # Main loop: will keep going until s == 0.
  
      while s and j < maxheight:
          # Choose a direction. -1 = backwards, 1 = forwards.
          v = int(2 * (np.random.uniform() < 0.5) - 1)
  
          # Double the size of the tree.
          if v == -1:
              (
                  thetaminus,
                  rminus,
                  gradminus,
                  _,
                  _,
                  _,
                  thetaprime,
                  gradprime,
                  logpprime,
                  extraprime,
                  rprime,
                  sprime,
                  alphaprime,
                  betaprime,
                  nalphaprime,
              ) = build_tree(
                  thetaminus,
                  rminus,
                  gradminus,
                  v,
                  j,
                  epsilon,
                  invmassmatrix,
                  f,
                  joint,
              )
          else:
              (
                  _,
                  _,
                  _,
                  thetaplus,
                  rplus,
                  gradplus,
                  thetaprime,
                  gradprime,
                  logpprime,
                  extraprime,
                  rprime,
                  sprime,
                  alphaprime,
                  betaprime,
                  nalphaprime,
              ) = build_tree(
                  thetaplus,
                  rplus,
                  gradplus,
                  v,
                  j,
                  epsilon,
                  invmassmatrix,
                  f,
                  joint,
              )
  
          assert beta > 0, beta
          assert betaprime >= 0, betaprime
  
          # Use Metropolis-Hastings to decide whether or not to move to a
          # point from the half-tree we just generated.
          if sprime and np.random.uniform() < betaprime / (
              beta + betaprime
          ):
              logp = logpprime
              grad = gradprime[:]
              theta = thetaprime
              extra = extraprime
              r0 = rprime
              # print("accepting", theta, logp)
  
          alpha += alphaprime
          beta += betaprime
          nalpha += nalphaprime
  
          # Decide if it's time to stop.
          sturn = stop_criterion(thetaminus, thetaplus, rminus, rplus)
          # print(sprime, sturn)
          s = sprime and sturn
          # Increment depth.
          j += 1
      # print("jumping to:", theta)
      # print('Tree height: %d, acceptance fraction: %03.2f%%/%03.2f%%, epsilon=%g' % (j, alpha/nalpha*100, beta/nalpha*100, epsilon))
      return alpha, beta, nalpha, theta, grad, logp, extra, r0, j
  
imports: import numpy
  import matplotlib.pyplot
  import scipy.special
  import scipy.stats
label_desc: Use image node_id 4 for calling a global function with example usage: tree_sample(theta, logp, r0, grad, extra, epsilon, invmassmatrix, f, joint, maxheight) and returns: alpha,  beta,  nalpha,  theta,  grad,  logp,  extra,  r0,  j