optimizer: jitting the optimizer step

Many optimizers mainly use element-wise operations in the update computation. This provides a good opportunity to `jit` this part of the computation to gain some perf benefits. Request here is to support jitting the optimizer step. This will also be a step towards being able to capture the complete computation in a training iteration i.e. forward pass, backward pass and parameters update.

A thing to note is that the optimizers provide by PyTorch all use inplace operations which are currently unsupported.
I can think of two potential approaches:
* One approach could be to write our own optimizer which is thunder friendly. We can use `prims.copy_` manually to implement the inplace update.
* Other would be to support all the required inplace operations in `torch/__init__.py` as decomposition of the original operation followed by `prims.copy_` into the input tensor.

Scope:
* [ ] Provide an implementation for Adam and AdamW optimizer such that we can use it in training lit-gpt model.