Need to access tape-producing and tape-consuming in std::autodiff for reusing a forward pass

I have a use case where I need to run reverse-mode autodiff more than once from a shared forward computation. Some intermediate computations are shared by multiple downstream quantities, so recomputing the entire forward pass for every reverse pass is wasteful.

For instance, one might have a shared forward computation that is computationally expensive, and rather than just loss_a and loss_b, one might have many more, i.e. loss_c, loss_d, ..., etc. that depend on the shared forward computation. Without retaining the computation graph, recomputing that shared forward pass for each loss would be unnecessarily expensive in both time and memory

In PyTorch via tch, the way to retain the computation graph is done by setting third argument of Tensor::run_backward to true:

use tch::Tensor;

// `parameters` must have `requires_grad` enabled.
let y = shared_forward(&parameters, &input);
let loss_a = objective_a(&y);
let loss_b = objective_b(&y);

// Keep the graph after this pass so it can be traversed again.
let grad_a = Tensor::run_backward(&[&loss_a], &[&parameters], true, false)
    .remove(0);

// This final pass can free the graph.
let grad_b = Tensor::run_backward(&[&loss_b], &[&parameters], false, false)
    .remove(0);

For completeness, this is the corresponding implementation of Tensor:: run_backward from https://docs.rs/tch/latest/tch/struct.Tensor.html#method.run_backward

 pub fn run_backward<T1, T2>(
        tensors: &[T1],
        inputs: &[T2],
        keep_graph: bool,
        create_graph: bool,
    ) -> Vec<Tensor>
    where
        T1: Borrow<Tensor>,
        T2: Borrow<Tensor>,
    {
        Tensor::f_run_backward(tensors, inputs, keep_graph, create_graph).unwrap()
    }

The Rust autodiff interface I would like is conceptually similar: a safe way to produce an opaque tape from a differentiated forward pass, then consume or reuse that tape for one or more reverse passes.

Illustrative pseudocode only—the exact API is not important:

let (output, tape) = shared_forward_augmented(&parameters, &input);

let grad_a = shared_forward_reverse(
    tape.borrow(),
    output_seed_for_objective_a,
)?;

let grad_b = shared_forward_reverse(
    tape,
    output_seed_for_objective_b,
)?;

The capability I would like from Rust autodiff is analogous: evaluate the shared forward computation once, retain its opaque reverse-mode state, and safely run multiple reverse passes for scalar outputs that depend on that shared computation.

Related thread: https://github.com/rust-lang/rust/issues/160806

Start by reading the existing std::autodiff interface and the related thread in issue 160806 to understand the current tape-producing and tape-consuming capabilities. The work is done when Rust has an agreed safe interface for retaining opaque reverse-mode state and running multiple reverse passes from one shared forward computation.