Repository: ivy-llc/ivy

Issue:
ivy.AdamW



Discussion:

Resolving patch:
diff --git a/ivy/stateful/optimizers.py b/ivy/stateful/optimizers.py
--- a/ivy/stateful/optimizers.py
+++ b/ivy/stateful/optimizers.py
@@ -417,6 +417,85 @@ def state(self):
         return ivy.Container({"mw": self._mw, "vw": self._vw})
 
 
+class AdamW(Adam):
+    def __init__(
+        self,
+        lr: float = 1e-4,
+        beta1: float = 0.9,
+        beta2: float = 0.999,
+        epsilon: float = 1e-07,
+        weight_decay: float = 0.0,
+        inplace: bool = True,
+        stop_gradients: bool = True,
+        trace_on_next_step: bool = False,
+        device: Optional[Union[ivy.Device, ivy.NativeDevice]] = None,
+    ):
+        """
+        Construct an ADAMW optimizer.
+
+        Parameters
+        ----------
+        lr
+            Learning rate, default is ``1e-4``.
+        beta1
+            gradient forgetting factor, default is ``0.9``
+        beta2
+            second moment of gradient forgetting factor, default is ``0.999``
+        epsilon
+            divisor during adamw update, preventing division by zero,
+            default is ``1e-07``
+        weight_decay
+            weight decay coefficient, default is ``0.0``
+        inplace
+            Whether to update the variables in-place, or to create new variable handles.
+            This is only relevant for frameworks with stateful variables such as
+            PyTorch.
+            Default is ``True``, provided the backend framework supports it.
+        stop_gradients
+            Whether to stop the gradients of the variables after each gradient step.
+            Default is ``True``.
+        trace_on_next_step
+            Whether to trace the optimizer on the next step. Default is ``False``.
+        device
+            Device on which to create the layer's variables 'cuda:0', 'cuda:1', 'cpu'
+            etc. (Default value = None)
+        """
+        self._weight_decay = weight_decay
+        super().__init__(
+            lr,
+            beta1,
+            beta2,
+            epsilon,
+            inplace,
+            stop_gradients,
+            trace_on_next_step,
+            device,
+        )
+
+    def _step(self, v: ivy.Container, grads: ivy.Container):
+        """
+        Update nested variables container v by AdamW update step, using nested grads
+        container.
+
+        Parameters
+        ----------
+        v
+            Nested variables to update.
+        grads
+            Nested gradients to update.
+
+        Returns
+        -------
+        ret
+            The updated variables, following AdamW update step.
+        """
+        # Apply L2 regularization directly to the parameters
+        if self._weight_decay != 0:
+            grads += self._weight_decay * v
+
+        return super()._step(v, grads)
+
+
 class LAMB(Optimizer):
     def __init__(
         self,


Test patch:
diff --git a/ivy_tests/test_ivy/test_stateful/test_optimizers.py b/ivy_tests/test_ivy/test_stateful/test_optimizers.py
--- a/ivy_tests/test_ivy/test_stateful/test_optimizers.py
+++ b/ivy_tests/test_ivy/test_stateful/test_optimizers.py
@@ -75,6 +75,73 @@ def test_adam_optimizer(
     )
 
 
+# AdamW
+@handle_method(
+    method_tree="AdamW._step",
+    dtype_x_lr=get_gradient_arguments_with_lr(
+        min_value=1e-05,
+        max_value=1e08,
+        num_arrays=2,
+        float_lr=True,
+        large_abs_safety_factor=2,
+        small_abs_safety_factor=2,
+    ),
+    beta1_n_beta2_n_epsilon=helpers.list_of_size(
+        x=helpers.floats(min_value=1e-1, max_value=1),
+        size=3,
+    ),
+    weight_decay=helpers.floats(min_value=0, max_value=1e-1),
+    inplace=st.booleans(),
+    stop_gradients=st.booleans(),
+    test_gradients=st.just(True),
+)
+def test_adamw_optimizer(
+    dtype_x_lr,
+    beta1_n_beta2_n_epsilon,
+    weight_decay,
+    inplace,
+    stop_gradients,
+    on_device,
+    class_name,
+    method_name,
+    backend_fw,
+    ground_truth_backend,
+    test_gradients,
+    init_flags,
+    method_flags,
+):
+    input_dtype, x, lr = dtype_x_lr
+    beta1, beta2, epsilon = beta1_n_beta2_n_epsilon
+    xs_grad_idxs = [[0, 0]] if method_flags.num_positional_args else [[1, "v"]]
+    helpers.test_method(
+        backend_to_test=backend_fw,
+        ground_truth_backend=ground_truth_backend,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        init_all_as_kwargs_np={
+            "lr": lr,
+            "beta1": beta1,
+            "beta2": beta2,
+            "epsilon": epsilon,
+            "weight_decay": weight_decay,
+            "inplace": inplace,
+            "stop_gradients": stop_gradients,
+        },
+        method_input_dtypes=input_dtype,
+        method_all_as_kwargs_np={
+            "v": x[0],
+            "grads": x[1],
+        },
+        class_name=class_name,
+        method_name=method_name,
+        rtol_=1e-1,
+        atol_=1e-1,
+        test_gradients=test_gradients,
+        xs_grad_idxs=xs_grad_idxs,
+        on_device=on_device,
+    )
+
+
 # lamb
 @handle_method(
     method_tree="LAMB._step",