human
2 Define DoSE helper classes and functions
assistant
#Here is the Python program you requested.
def get_DoSE_KDE(T, dataset):
  """Get a distribution and decision rule for OOD detection using DoSE.

  Given a tensor of statistics tx, compute a Kernel Density Estimate (KDE) of
  the statistics.  This uses a quantiles trick to cut down the number of 
  samples used in the KDE (to lower the cost of evaluating a trial point).

  Args:
    T:  A function which takes an input image and returns a vector of
      statistics evaluated using the model.
    dataset:  A tensorflow_datasets `Dataset` which will be used to evaluate
      statistics to construct the estimator.
  Returns:
    is_ood:  A function which takes a new point `x` and `threshold`, and 
      computes the decision rule KDE.log_prob(T(x)) < threshold
    dose_kde:  A tfd.MixtureSameFamily object.  The distribution used as the KDE
      from which the log_prob of a batch of statistics can be computed.
  """

  # First we should evaluate the statistics on the training set.
  it = iter(dataset)
  for x, y in it:
    if not "tx" in locals():
      tx = T(x)
    else:
      tx = tf.concat([tx, T(x)], axis=-1)

  n = tf.cast(tf.shape(tx)[-1], tx.dtype)
  num_quantiles = int(25)
  q = tfp.stats.quantiles(tx, num_quantiles, axis=-1)
  q = tf.transpose(q, tf.roll(tf.range(tf.rank(q)), shift=-1, axis=0))
  # Scott's Rule:
  h = 3.49 * tf.math.reduce_std(tx, axis=-1, keepdims=True) * (n)**(-1./3.)
  h *= n / num_quantiles
 
  dose_kde = tfd.Independent(
      tfd.MixtureSameFamily(
          mixture_distribution=tfd.Categorical(logits=tf.zeros(num_quantiles + 1)),
          components_distribution=tfd.Normal(loc=q, scale=h)),
      reinterpreted_batch_ndims=1)
  is_ood = lambda x, threshold: dose_kde.log_prob(tf.transpose(T(x), [1, 0])) < tf.math.log(threshold)
  dose_log_prob = lambda x: dose_kde.log_prob(tf.transpose(T(x), [1, 0]))

  # T(x) returns shape [T, N], but dose_kde works on shape [N, T]
  return is_ood, dose_kde, dose_log_prob, tx

class DoSE_administrator(object):
  def __init__(self, T, train_dataset, eval_dataset):
    """Administrate DoSE...