feat: add build

build/torch26-cxx11-cu118-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch26-cxx11-cu118-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3bf69bbdc95e26ae0b98ecd42e238a8fa67c503348ba062c8af18e681b758db3
+size 2900352

build/torch26-cxx11-cu118-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"

build/torch26-cxx11-cu124-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch26-cxx11-cu124-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50887573fa0599bcc94b948faaa38d9c6e06f8a654c066d5f49460b86a109c1b
+size 2929048

build/torch26-cxx11-cu124-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"

build/torch26-cxx11-cu126-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch26-cxx11-cu126-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eec0b7d37568f183dfbcb419f359f9b046fd893e075525083f635c2f936c89e0
+size 2933584

build/torch26-cxx11-cu126-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"

build/torch26-cxx98-cu118-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch26-cxx98-cu118-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13a70e896d05f084a9a736b17604e859a6b384fb9dc93b5c28486a3a84d2bc93
+size 2897504

build/torch26-cxx98-cu118-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"

build/torch26-cxx98-cu124-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch26-cxx98-cu124-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f11e153608ac8a325a312278994b01278dd47c41173ee06241eff26f69637a48
+size 2922152

build/torch26-cxx98-cu124-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"

build/torch26-cxx98-cu126-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch26-cxx98-cu126-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:993bdb2c8dd5dc103bdb1c2d632e0f41ace2caf6665e3211b2954bc191eb5bf9
+size 2926688

build/torch26-cxx98-cu126-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"

build/torch27-cxx11-cu118-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch27-cxx11-cu118-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47610d48450a101312d6e6335153e925a90efd357e1b252f3eec07b1459cd58f
+size 2900448

build/torch27-cxx11-cu118-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"

build/torch27-cxx11-cu126-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch27-cxx11-cu126-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d202c842a41a288e7c56b9063ad1ab1962cde09a647072344f07c76687555f7
+size 2933616

build/torch27-cxx11-cu126-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"

build/torch27-cxx11-cu128-x86_64-linux/adam_atan2/__init__.py

@@ -0,0 +1,133 @@
+# NOTE: Torch needs to be imported before the custom
+# extensions. Otherwise libc10.so cannot be found.
+import torch
+from ._ops import ops
+
+from typing import List, Tuple, Union
+from torch import Tensor
+from torch.optim.optimizer import Optimizer, ParamsT
+
+
+class AdamATan2(Optimizer):
+    def __init__(
+        self,
+        params: ParamsT,
+        lr: Union[float, Tensor] = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        weight_decay: float = 1e-2,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+
+        defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+    ):
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            params_with_grad.append(p)
+            if p.grad.is_sparse:
+                raise RuntimeError("AdamW does not support sparse gradients")
+            grads.append(p.grad)
+
+            state = self.state[p]
+
+            # State initialization
+            if len(state) == 0:
+                # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off.
+                # This is because kernel launches are costly on CUDA and XLA.
+                state["step"] = torch.zeros((), dtype=torch.float32, device=p.device)
+                # Exponential moving average of gradient values
+                state["exp_avg"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+                # Exponential moving average of squared gradient values
+                state["exp_avg_sq"] = torch.zeros_like(
+                    p, memory_format=torch.preserve_format
+                )
+
+            exp_avgs.append(state["exp_avg"])
+            exp_avg_sqs.append(state["exp_avg_sq"])
+            state_steps.append(state["step"])
+
+    def step(self):
+        """Perform a single optimization step."""
+        self._cuda_graph_capture_health_check()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            state_steps = []
+            beta1, beta2 = group["betas"]
+
+            self._init_group(
+                group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps
+            )
+
+            _adam_atan2(
+                params_with_grad,
+                grads,
+                exp_avgs,
+                exp_avg_sqs,
+                state_steps,
+                beta1=beta1,
+                beta2=beta2,
+                lr=group["lr"],
+                weight_decay=group["weight_decay"],
+            )
+
+
+def _adam_atan2(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    state_steps: List[Tensor],
+    beta1: float,
+    beta2: float,
+    lr: float,
+    weight_decay: float,
+) -> None:
+    if not params:
+        return
+
+    # We only support scalar lr.
+    assert not isinstance(lr, Tensor)
+
+    grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(
+        [params, grads, exp_avgs, exp_avg_sqs, state_steps]
+    )
+    for (device, _), (
+        (
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+        ),
+        _,
+    ) in grouped_tensors.items():
+        torch._foreach_add_(device_state_steps, 1)
+        ops.adam_atan2_cuda_impl_(
+            device_params,
+            device_grads,
+            device_exp_avgs,
+            device_exp_avg_sqs,
+            device_state_steps,
+            lr,
+            beta1,
+            beta2,
+            weight_decay,
+        )

build/torch27-cxx11-cu128-x86_64-linux/adam_atan2/_adam_atan2_40f2269.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:37cf6b90c834d55ed802734bb91a8a601f6eab6a88e0e8eed7bd4cb449c563fd
+size 3688960

build/torch27-cxx11-cu128-x86_64-linux/adam_atan2/_ops.py

@@ -0,0 +1,9 @@
+import torch
+from . import _adam_atan2_40f2269
+ops = torch.ops._adam_atan2_40f2269
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_adam_atan2_40f2269::{op_name}"