cover forward + loss + backward

## 🚀 Feature

Add a way to not only compile the model, but the `forward` + `loss` + `loss.backward` (later also `optimizer.step` + `optimizer.zero_grad`, but not here).

### Motivation

For training applications, we currently integrate with torch autograd, but we could and should aim to extend beyond this. Also, this would simplify a things, e.g. there would be no need to split the trace into forward and backward (but we need to figure out re-computation nonetheless).

### Pitch

The UX might look like this:

```python
def fn(model, inp, target):
out = model(inp)
loss = lossfn(out, target)
loss.backward()

jfn = jit(fn, models=(model,))
```

### Alternatives

Other UX possible.

### Implementation thoughts

Maybe we should figure out how to combine prologues in general. This would have applications here (how do we differentiate the model prologue from the entire function prologue) but also elsewhere.

Start by reading the existing jit API and its torch autograd integration, then examine how the current trace is split between forward and backward and how prologues are handled. Done means a function containing model forward, loss computation, and loss.backward can be compiled through jit without requiring separate traces; optimizer steps are explicitly out of scope.

Written by the indexing model from the issue text.