How to create an AdamW optimizer with Amsgrad

[SalamanderXing]

Hi, I'm reproducing a deep learning paper from PyTorch -> JAX. The paper uses the following optimizer:

optimizer = torch.optim.AdamW(
       model.parameters(),
       lr=lr,
       amsgrad=True,
       weight_decay=weight_decay,
)

I found both AdamW and AMSGrad in Optax, but not AMSGrad with weight decay