R/inference.R

In greta: Simple and 'Scalable' Statistical Modelling in R

#' @name inference
#' @title Statistical inference on greta models.
#' @description Carry out statistical inference on greta models by
#'   MCMC or likelihood/posterior optimisation.
NULL


# create an object stash in greta's namespace, to return traces to the user when
# they abort a run
greta_stash <- new.env()
greta_stash$python_has_been_initialised <- FALSE
greta_stash$deps_removed_this_session <- FALSE
greta_stash$numerical_messages <- c(
  "is not invertible",
  "Cholesky decomposition was not successful"
)

#' @rdname inference
#' @export
#' @importFrom stats rnorm runif
#' @importFrom utils setTxtProgressBar txtProgressBar
#' @importFrom future plan nbrOfWorkers
#'
#' @param model greta_model object
#' @param sampler sampler used to draw values in MCMC. See
#'   [samplers()] for options.
#' @param n_samples number of MCMC samples to draw per chain (after any warm-up,
#'   but before thinning)
#' @param thin MCMC thinning rate; every `thin` samples is retained, the
#'   rest are discarded
#' @param warmup number of samples to spend warming up the mcmc sampler (moving
#'   chains toward the highest density area and tuning sampler hyperparameters).
#' @param chains number of MCMC chains to run. Default is 2. We recommend using more chains as this helps improve convergence. However the number of chains specified can increase the CPU load, so we have to set a lower default value.
#' @param n_cores the maximum number of CPU cores used by each sampler (see
#'   details). If NULL (default), it sets them to 2 cores.
#' @param verbose whether to print progress information to the console
#' @param pb_update how regularly to update the progress bar (in iterations).
#'   If `pb_update` is less than or equal to `thin`, it will be set
#'   to `thin + 1` to ensure at least one saved iteration per
#'   `pb_update` iterations.
#' @param one_by_one whether to run TensorFlow MCMC code one iteration at a
#'   time, so that greta can handle numerical errors as 'bad' proposals (see
#'   below).
#' @param initial_values an optional `initials` object (or list of
#'   `initials` objects of length `chains`) giving initial values for
#'   some or all of the variables in the model. These will be used as the
#'   starting point for sampling/optimisation.
#' @param trace_batch_size the number of posterior samples to process at a time
#'   when tracing the parameters of interest; reduce this to reduce memory
#'   demands
#'
#' @param compute_options Default is to use CPU only with `cpu_only()`. Use
#'   `gpu_only()` to use only GPU. In the future we will add more options for
#'   specifying CPU and GPU use.
#'
#' @details For `mcmc()` if `verbose = TRUE`, the progress bar shows
#'   the number of iterations so far and the expected time to complete the phase
#'   of model fitting (warmup or sampling). Occasionally, a proposed set of
#'   parameters can cause numerical instability (I.e. the log density or its
#'   gradient is `NA`, `Inf` or `-Inf`); normally because the log
#'   joint density is so low that it can't be represented as a floating point
#'   number. When this happens, the progress bar will also display the
#'   proportion of proposals so far that were 'bad' (numerically unstable) and
#'   therefore rejected. Some numerical instability during the warmup phase is
#'   normal, but 'bad' samples during the sampling phase can lead to bias in
#'   your posterior sample. If you only have a few bad samples (<10\\%), you can
#'   usually resolve this with a longer warmup period or by manually defining
#'   starting values to move the sampler into a more reasonable part of the
#'   parameter space. If you have more samples than that, it may be that your
#'   model is misspecified. You can often diagnose this by using
#'   [calculate()] to evaluate the values of greta arrays, given
#'   fixed values of model parameters, and checking the results are what you
#'   expect.
#'
#'   greta runs multiple chains simultaneously with a single sampler,
#'   vectorising all operations across the chains. E.g. a scalar addition in
#'   your model is computed as an elementwise vector addition (with vectors
#'   having length `chains`), a vector addition is computed as a matrix
#'   addition etc. TensorFlow is able to parallelise these operations, and this
#'   approach reduced computational overheads, so this is the most efficient of
#'   computing on multiple chains.
#'
#'   Multiple mcmc samplers (each of which can simultaneously run multiple
#'   chains) can also be run in parallel by setting the execution plan with the
#'   `future` package. Only `plan(multisession)` futures or
#'   `plan(cluster)` futures that don't use fork clusters are allowed,
#'   since forked processes conflict with TensorFlow's parallelism. Explicitly
#'   parallelising chains on a local machine with `plan(multisession)` will
#'   probably be slower than running multiple chains simultaneously in a single
#'   sampler (with `plan(sequential)`, the default) because of the overhead
#'   required to start new sessions. However, `plan(cluster)` can be used
#'   to run chains on a cluster of machines on a local or remote network. See
#'   [future::cluster()] for details, and the
#'   `future.batchtools` package to set up plans on clusters with job
#'   schedulers.
#'
#'   If `n_cores = NULL` and mcmc samplers are being run sequentially, each
#'   sampler will be allowed to use only 2 CPU cores (possibly to compute
#'   multiple chains sequentially). If samplers are being run in parallel with
#'   the `future` package, `n_cores` will be set so that `n_cores *
#'   [future::nbrOfWorkers]` is less than the number
#'   of CPU cores.
#'
#'   After carrying out mcmc on all the model parameters, `mcmc()`
#'   calculates the values of (i.e. traces) the parameters of interest for each
#'   of these samples, similarly to [calculate()]. Multiple
#'   posterior samples can be traced simultaneously, though this can require
#'   large amounts of memory for large models. As in `calculate`, the
#'   argument `trace_batch_size` can be modified to trade-off speed against
#'   memory usage.
#'
#' @note to set a seed with MCMC you can use [set.seed()], or
#'   [tensorflow::set_random_seed()]. They both give identical results. See
#'   examples below.
#'
#' @return `mcmc`, `stashed_samples` & `extra_samples` - a
#'   `greta_mcmc_list` object that can be analysed using functions from the
#'   coda package. This will contain mcmc samples of the greta arrays used to
#'   create `model`.
#'
#' @examples
#' \dontrun{
#' # define a simple Bayesian model
#' x <- rnorm(10)
#' mu <- normal(0, 5)
#' sigma <- lognormal(1, 0.1)
#' distribution(x) <- normal(mu, sigma)
#' m <- model(mu, sigma)
#'
#' # carry out mcmc on the model
#' draws <- mcmc(m, n_samples = 100)
#'
#' # add some more samples
#' draws <- extra_samples(draws, 200)
#'
#' # initial values can be passed for some or all model variables
#' draws <- mcmc(m, chains = 1, initial_values = initials(mu = -1))
#'
#' # if there are multiple chains, a list of initial values should be passed,
#' # othewise the same initial values will be used for all chains
#' inits <- list(initials(sigma = 0.5), initials(sigma = 1))
#' draws <- mcmc(m, chains = 2, initial_values = inits)
#'
#' # you can auto-generate a list of initials with something like this:
#' inits <- replicate(4,
#'   initials(mu = rnorm(1), sigma = runif(1)),
#'   simplify = FALSE
#' )
#' draws <- mcmc(m, chains = 4, initial_values = inits)
#'
#' # or find the MAP estimate
#' opt_res <- opt(m)
#'
#' # get the MLE of the normal variance
#' mu <- variable()
#' variance <- variable(lower = 0)
#' distribution(x) <- normal(mu, sqrt(variance))
#' m2 <- model(variance)
#'
#' # adjust = FALSE skips the jacobian adjustments used in MAP estimation, to
#' # give the true maximum likelihood estimates
#' o <- opt(m2, adjust = FALSE)
#'
#' # the MLE corresponds to the *unadjusted* sample variance, but differs
#' # from the sample variance
#' o$par
#' mean((x - mean(x))^2) # same
#' var(x) # different
#'
#' # initial values can also be passed to optimisers:
#' o <- opt(m2, initial_values = initials(variance = 1))
#'
#' # and you can return a list of the Hessians for each of these parameters
#' o <- opt(m2, hessian = TRUE)
#' o$hessian
#'
#'
#' # to get a hessian matrix across multiple greta arrays, you must first
#' # combine them and then split them up for use in the model (so that the
#' # combined vector is part of the model) and pass that vector to model:
#' params <- c(variable(), variable(lower = 0))
#' mu <- params[1]
#' variance <- params[2]
#' distribution(x) <- normal(mu, sqrt(variance))
#' m3 <- model(params)
#' o <- opt(m3, hessian = TRUE)
#' o$hessian
#'
#' # using set.seed or tensorflow::set_random_seed to set RNG for MCMC
#' a <- normal(0, 1)
#' y <- normal(a, 1)
#' m <- model(y)
#'
#' set.seed(12345)
#' one <- mcmc(m, n_samples = 1, chains = 1)
#' set.seed(12345)
#' two <- mcmc(m, n_samples = 1, chains = 1)
#' # same
#' all.equal(as.numeric(one), as.numeric(two))
#' tensorflow::set_random_seed(12345)
#' one_tf <- mcmc(m, n_samples = 1, chains = 1)
#' tensorflow::set_random_seed(12345)
#' two_tf <- mcmc(m, n_samples = 1, chains = 1)
#' # same
#' all.equal(as.numeric(one_tf), as.numeric(two_tf))
#' # different
#' all.equal(as.numeric(one), as.numeric(one_tf))
#'
#' }
mcmc <- function(
  model,
  sampler = hmc(),
  n_samples = 1000,
  thin = 1,
  warmup = 1000,
  chains = 2,
  n_cores = NULL,
  verbose = TRUE,
  pb_update = 50,
  one_by_one = FALSE,
  initial_values = initials(),
  trace_batch_size = 100,
  compute_options = cpu_only()
) {
  # set device to be CPU/GPU for the entire run
  with(tf$device(compute_options), {
    # message users about random seeds and GPU usage if they are using GPU
    message_if_using_gpu(compute_options)

    check_if_greta_array_in_mcmc(model)

    check_trace_batch_size(trace_batch_size)

    check_not_data_greta_arrays(model)

    chains <- check_positive_integer(chains, "chains")

    # get the dag containing the target nodes
    dag <- model$dag

    # turn initial values into a list if needed (checking length). Convert from
    # a named list on the constrained scale to free state vectors
    initial_values <- prep_initials(initial_values, chains, dag)

    samplers <- build_samplers(
      sampler = sampler,
      initial_values = initial_values,
      chains = chains,
      model = model,
      compute_options = compute_options
    )

    # if verbose = FALSE, make pb_update as big as possible to speed up sampling
    if (!verbose) {
      pb_update <- Inf
    }

    # now make it finite
    pb_update <- min(pb_update, max(warmup, n_samples))
    pb_update <- max(pb_update, thin + 1)

    run_samplers(
      samplers = samplers,
      n_samples = n_samples,
      thin = thin,
      warmup = warmup,
      verbose = verbose,
      pb_update = pb_update,
      one_by_one = one_by_one,
      n_cores = n_cores,
      from_scratch = TRUE,
      trace_batch_size = trace_batch_size,
      compute_options = compute_options
    )

    # close `with` setting to control CPU/GPU usage
  })
}

#' @importFrom future future resolved value
run_samplers <- function(
  samplers,
  n_samples,
  thin,
  warmup,
  verbose,
  pb_update,
  one_by_one,
  n_cores,
  from_scratch,
  trace_batch_size,
  compute_options
) {
  # check the future plan is valid, and get information about it
  plan_is <- check_future_plan()

  # coerce the iterations to integer
  n_samples <- as.integer(n_samples)
  warmup <- as.integer(warmup)
  thin <- as.integer(thin)

  dag <- samplers[[1]]$model$dag
  chains <- samplers[[1]]$n_chains
  n_cores <- check_n_cores(n_cores, length(samplers), plan_is)
  float_type <- dag$tf_float

  # stash the samplers now, to retrieve draws later
  greta_stash$samplers <- samplers

  inform_if_local_parallel_multiple_samplers(
    plan_is,
    samplers,
    n_cores,
    compute_options
  )

  inform_if_remote_machine(plan_is, samplers)

  n_chain <- length(samplers)
  chains <- seq_len(n_chain)

  # determine the type of progress information
  set_progress_bar_type(n_chain)

  parallel_reporting <- verbose &
    plan_is$parallel &
    plan_is$local &
    !is.null(greta_stash$callbacks)

  # if we're running in parallel and there are callbacks registered,
  # give the samplers somewhere to write their progress
  if (parallel_reporting) {
    sampler <- sampler_parallel_reporting(
      n_chain,
      samplers,
      chains,
      n_samples,
      warmup
    )
  }

  if (plan_is$parallel) {
    dispatch <- future::future
  } else {
    dispatch <- function(expr, ...) expr
  }

  # dispatch all the jobs
  for (chain in chains) {
    sampler <- samplers[[chain]]
    sampler$compute_options <- compute_options
    samplers[[chain]] <- dispatch(
      sampler$run_chain(
        n_samples = n_samples,
        thin = thin,
        warmup = warmup,
        verbose = verbose,
        pb_update = pb_update,
        one_by_one = one_by_one,
        plan_is = plan_is,
        n_cores = n_cores,
        float_type = float_type,
        trace_batch_size = trace_batch_size,
        from_scratch = from_scratch
      ),
      seed = future_seed()
    )
  }

  # if we're non-sequential and there's a callback registered,
  # loop until they are resolved, executing the callbacks
  if (parallel_reporting) {
    while (!all(vapply(samplers, resolved, FALSE))) {
      # loop through callbacks executing them
      for (callback in greta_stash$callbacks) {
        callback()
      }

      # get some nap time
      Sys.sleep(0.1)
    }

    cat("\n")
  }

  # if we were running in parallel, retrieve the samplers and put them back in
  # the stash to return
  if (plan_is$parallel) {
    samplers <- lapply(samplers, future::value)
    greta_stash$samplers <- samplers
  }

  # get chains from the samplers, with raw values as an attribute
  draws <- stashed_samples()

  # empty the stash
  rm("samplers", envir = greta_stash)

  draws
}

#' @rdname inference
#' @export
#' @importFrom stats na.omit
#' @importFrom coda mcmc.list
#'
#' @details If the sampler is aborted before finishing (and `future`
#'   parallelism isn't being used), the samples collected so far can be
#'   retrieved with `stashed_samples()`. Only samples from the sampling
#'   phase will be returned.
stashed_samples <- function() {
  stashed <- exists("samplers", envir = greta_stash)

  if (stashed) {
    samplers <- greta_stash$samplers

    trace_names <- samplers[[1]]$model$dag$trace_names

    # get draws as 3D arrays for each sampler
    free_state_draws <- lapply(samplers, member, "traced_free_state")
    values_draws <- lapply(samplers, member, "traced_values")

    # split along chain dimension, to get a list of chains
    free_state_draws <- unlist(free_state_draws, recursive = FALSE)
    values_draws <- unlist(values_draws, recursive = FALSE)

    # apply the trace names back to them
    values_draws <- lapply(values_draws, `colnames<-`, trace_names)

    # if there are no samples, return a list of NULLs
    no_samples <- nrow(values_draws[[1]]) == 0
    if (no_samples) {
      return(replicate(length(samplers), NULL))
    } else {
      thins <- lapply(samplers, member, "thin")

      # convert to mcmc objects, passing on thinning
      free_state_draws <- mapply(
        prepare_draws,
        draws = free_state_draws,
        thin = thins,
        SIMPLIFY = FALSE
      )
      values_draws <- mapply(
        prepare_draws,
        draws = values_draws,
        thin = thins,
        SIMPLIFY = FALSE
      )

      # convert to mcmc.list objects
      free_state_draws <- coda::mcmc.list(free_state_draws)
      values_draws <- coda::mcmc.list(values_draws)

      # prep the raw model objects
      model_info <- list(
        raw_draws = free_state_draws,
        samplers = samplers,
        model = samplers[[1]]$model,
        warmup = samplers[[1]]$warmup
      )

      values_draws <- as_greta_mcmc_list(values_draws, model_info)

      return(values_draws)
    }
  } else {
    return(invisible())
  }
}

#' @rdname inference
#'
#' @export
#'
#' @param draws a greta_mcmc_list object returned by `mcmc` or
#'   `stashed_samples`
#'
#' @details Samples returned by `mcmc()` and `stashed_samples()` can
#'   be added to with `extra_samples()`. This continues the chain from the
#'   last value of the previous chain and uses the same sampler and model as was
#'   used to generate the previous samples. It is not possible to change the
#'   sampler or extend the warmup period.
#'
extra_samples <- function(
  draws,
  n_samples = 1000,
  thin = 1,
  n_cores = NULL,
  verbose = TRUE,
  pb_update = 50,
  one_by_one = FALSE,
  trace_batch_size = 100,
  compute_options = cpu_only()
) {
  model_info <- get_model_info(draws)
  samplers <- model_info$samplers

  # set the last values as the current free state values
  for (sampler in samplers) {
    free_state_draws <- sampler$traced_free_state
    n_draws <- nrow(free_state_draws[[1]])
    free_state_draws <- lapply(free_state_draws, `[`, n_draws, )
    sampler$free_state <- do.call(rbind, free_state_draws)
  }

  run_samplers(
    samplers = samplers,
    n_samples = n_samples,
    thin = thin,
    warmup = 0L,
    verbose = verbose,
    pb_update = pb_update,
    one_by_one = one_by_one,
    n_cores = n_cores,
    from_scratch = FALSE,
    trace_batch_size = trace_batch_size,
    compute_options = compute_options
  )
}

# convert some 'data' values from the constrained to the free state, for a given
# 'node'
#' @importFrom stats qlogis
to_free <- function(node, data) {
  # use reverse mode of bijectors!

  lower <- node$lower
  upper <- node$upper

  # TODO
  # replace these with more informative errors related to the range of values
  unsupported_error <- function() {
    cli::cli_abort(
      "Some provided initial values are outside the range of values their \\
        variables can take"
    )
  }

  high <- function(x) {
    initials_outside_support <- any(x <= lower)
    if (initials_outside_support) {
      unsupported_error()
    }
    log(x - lower)
  }

  low <- function(x) {
    initials_outside_support <- any(x >= upper)
    if (initials_outside_support) {
      unsupported_error()
    }
    log(upper - x)
  }

  both <- function(x) {
    initials_outside_support <- any(x >= upper | x <= lower)
    if (initials_outside_support) {
      unsupported_error()
    }
    stats::qlogis((x - lower) / (upper - lower))
  }

  fun <- switch(
    node$constraint,
    scalar_all_none = identity,
    scalar_all_high = high,
    scalar_all_low = low,
    scalar_all_both = both
  )

  fun(data)
}

# convert a named list of initial values into the corresponding vector of values
# on the free state
parse_initial_values <- function(initials, dag) {
  # skip if no inits provided
  if (identical(initials, initials())) {
    free_parameters <- dag$example_parameters(free = TRUE)
    free_parameters <- unlist_tf(free_parameters)
    return(free_parameters)
  }

  # find the elements we have been given initial values for
  tf_names <- vapply(
    names(initials),
    function(name, env) {
      ga <- get(name, envir = env)
      node <- get_node(ga)
      dag$tf_name(node)
    },
    env = parent.frame(4),
    FUN.VALUE = ""
  )

  check_greta_arrays_associated_with_model(tf_names)

  params <- dag$example_parameters(free = FALSE)
  idx <- match(tf_names, names(params))

  # make nodes do this conversion and checking in the future also make them
  # handle more complex situations (like Wishart, which is an operation with a
  # distribution, so needs to find the corresponding variable), and check the
  # provided values are in the support of the (constrained scale of the)
  # variable

  # find the corresponding nodes and check they are variable nodes
  forward_names <- glue::glue("all_forward_{dag$node_tf_names}")
  nodes <- dag$node_list[match(tf_names, forward_names)]

  check_nodes_all_variable(nodes)

  target_dims <- lapply(params[idx], dim)
  replacement_dims <- lapply(initials, dim)
  check_initial_values_correct_dim(target_dims, replacement_dims)

  # convert the initial values to their free states
  inits_free <- mapply(to_free, nodes, initials, SIMPLIFY = FALSE)

  # set them in the list and flatten to a vector in the same order as tensorflow
  params[idx] <- inits_free
  params <- unlist_tf(params)

  # force them to be row vectors and return
  params <- matrix(params, nrow = 1)
  params
}

# convert (possibly NULL) user-specified initial values into a list of the
# correct length, with nice error messages
prep_initials <- function(initial_values, n_chains, dag) {
  # if the user passed a single set of initial values, repeat them for all
  # chains
  if (is.initials(initial_values)) {
    inform_if_one_set_of_initials(initial_values, n_chains)

    initial_values <- replicate(n_chains, initial_values, simplify = FALSE)
  }

  # TODO: revisit logic here for errors and messages
  check_initial_values_match_chains(initial_values, n_chains)
  check_initial_values_correct_class(initial_values)

  # convert them to free state vectors
  initial_values <- lapply(
    initial_values,
    parse_initial_values,
    dag
  )

  initial_values
}

#' @rdname inference
#' @export
#' @param ... named numeric values, giving initial values of some or all of the
#'   variables in the model (unnamed variables will be automatically
#'   initialised)
#'
initials <- function(...) {
  values <- list(...)
  names <- names(values)

  check_initials_are_named(values)

  # coerce to greta-array-like shape
  values <- lapply(values, as_2d_array)

  check_initials_are_numeric(values)

  class(values) <- c("initials", class(values))
  values
}

#' @export
print.initials <- function(x, ...) {
  empty_initials <- identical(x, initials())
  if (empty_initials) {
    cat("an empty greta initials object")
  } else {
    cat("a greta initials object with values:\n\n")
    print(unclass(x))
  }

  invisible(x)
}

#' @rdname inference
#' @export
#'
#' @param max_iterations the maximum number of iterations before giving up
#' @param tolerance the numerical tolerance for the solution, the optimiser
#'   stops when the (absolute) difference in the joint density between
#'   successive iterations drops below this level
#' @param optimiser an `optimiser` object giving the optimisation algorithm
#'   and parameters See [optimisers()].
#' @param adjust whether to account for Jacobian adjustments in the joint
#'   density. Set to `FALSE` (and do not use priors) for maximum likelihood
#'   estimates, or `TRUE` for maximum *a posteriori* estimates.
#' @param hessian whether to return a list of *analytically* differentiated
#'   Hessian arrays for the parameters
#' @param compute_options Default is to use CPU only with `cpu_only()`. Use
#'   `gpu_only()` to use only GPU. In the future we will add more options for
#'   specifying CPU and GPU use.
#'
#' @details Because `opt()` acts on a list of greta arrays with possibly
#'   varying dimension, the `par` and `hessian` objects returned by
#'   `opt()` are named lists, rather than a vector (`par`) and a
#'   matrix (`hessian`), as returned by [stats::optim()].
#'   Because greta arrays may not be vectors, the Hessians may not be matrices,
#'   but could be higher-dimensional arrays. To return a Hessian matrix covering
#'   multiple model parameters, you can construct your model so that all those
#'   parameters are in a vector, then split the vector up to define the model.
#'   The parameter vector can then be passed to model. See example.
#'

#' @return `opt` - a list containing the following named elements:
#'   \itemize{
#'    \item `par` a named list of the optimal values for the greta arrays
#'     specified in `model`
#'    \item `value` the (unadjusted) negative log joint density of the
#'     model at the parameters 'par'
#'    \item `iterations` the number of iterations taken by the optimiser
#'    \item `convergence` an integer code, 0 indicates successful
#'     completion, 1 indicates the iteration limit `max_iterations` had
#'     been reached
#'   \item `hessian` (if `hessian = TRUE`) a named list of hessian
#'     matrices/arrays for the parameters (w.r.t. `value`)
#'  }
#'
opt <- function(
  model,
  optimiser = bfgs(),
  max_iterations = 100,
  tolerance = 1e-6,
  initial_values = initials(),
  adjust = TRUE,
  hessian = FALSE,
  compute_options = cpu_only()
) {
  # set device to be CPU/GPU for the entire run
  with(tf$device(compute_options), {
    # message users about random seeds and GPU usage if they are using GPU
    message_if_using_gpu(compute_options)

    # check initial values. Can up the number of chains in the future to handle
    # random restarts
    initial_values_list <- prep_initials(initial_values, 1, model$dag)

    # create R6 object of the right type
    object <- optimiser$class$new(
      initial_values = initial_values_list[1],
      model = model,
      name = optimiser$name,
      method = optimiser$method,
      parameters = optimiser$parameters,
      other_args = optimiser$other_args,
      max_iterations = max_iterations,
      tolerance = tolerance,
      adjust = adjust
    )

    # run it and get the outputs
    object$run()
    outputs <- object$return_outputs()

    # optionally evaluate the hessians at these parameters (return as hessian for
    # objective function)
    if (hessian) {
      which_objective <- ifelse(adjust, "adjusted", "unadjusted")
      hessians <- model$dag$hessians(
        # coercing to Variable as TFP free state is a Tensor not a variable
        # and you can overload a Variable as a Variable without consequence
        # (hopefully)
        free_state = tf$Variable(object$free_state),
        which_objective = which_objective
      )
      hessians <- lapply(hessians, `*`, -1)
      outputs$hessian <- hessians
    }

    outputs

    # close setting of CPU/GPU usage
  })
}

inference_module <- module(dag_class, progress_bar = progress_bar_module)