Feature Request: Public API for compiled validation using `TrainState` cache

[msainsburydale]

When using TrainState with ReactantDevice, single_train_step! handles compilation and caching of the training step internally via TrainState.cache. However, there is no public API to perform a compiled forward pass for validation that shares this compilation cache. Is that possible?

The reason I ask is that a typical training loop requires both:

1. A training step — compute loss, compute gradients, and update parameters (single_train_step! handles this)
2. A validation step — compute loss without computing gradients.

However, with the current public API, I can't see how Step 2 can be done efficiently (compute_gradients uses the cache of TrainState but computes the gradients, which is wasteful). I know that @compile exists, but this doesn't share the cache with the TrainState object and doesn't include the user-friendly set-up of the public TrainState constructor.

Would it be hard to split compute_gradients into a function compute_objective, which does:

loss, stats, ts = compute_objective(objective_function, data, ts)