std::unordered_map<int64_t, int64_t> tensor_usage_;

  // If false, all activations are deleted in the first call to ComputeGradient.
  // Else, only when this is destructed.
  bool persistent_;
};

// Describes a callback for special-cased and more efficient jvp computation.
//
// Could just be a simple typedef in ForwardAccumulator, but MSVC chokes on
// that.
template <typename Gradient>
class ForwardFunction