From b28f3e82a7fddc23599dc07d49eb9d833879a5eb Mon Sep 17 00:00:00 2001
From: Mayank Mishra <32954280+mayank31398@users.noreply.github.com>
Date: Wed, 7 Aug 2024 00:50:38 -0400
Subject: [PATCH] rename module (#3)

---
 kernel_hyperdrive/utils.py                                | 2 +-
 kernel_hyperdrive/vector_addition/cuda_kernel/__init__.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/kernel_hyperdrive/utils.py b/kernel_hyperdrive/utils.py
index 0ec14b59..31584e0d 100644
--- a/kernel_hyperdrive/utils.py
+++ b/kernel_hyperdrive/utils.py
@@ -5,7 +5,7 @@
 
 def compile_helpers() -> None:
     load_cpp_extension(
-        "vector_addition_cuda",
+        "khd",
         sources=[
             os.path.join(os.path.dirname(__file__), "vector_addition/cuda_kernel/vector_addition.cpp"),
             os.path.join(os.path.dirname(__file__), "vector_addition/cuda_kernel/vector_addition.cu"),
diff --git a/kernel_hyperdrive/vector_addition/cuda_kernel/__init__.py b/kernel_hyperdrive/vector_addition/cuda_kernel/__init__.py
index 211184f9..828e590a 100644
--- a/kernel_hyperdrive/vector_addition/cuda_kernel/__init__.py
+++ b/kernel_hyperdrive/vector_addition/cuda_kernel/__init__.py
@@ -6,9 +6,9 @@
 
 class _VectorAddition_CUDA(torch.autograd.Function):
     def forward(ctx, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
-        import vector_addition_cuda
+        import khd
 
-        return vector_addition_cuda.vector_addition_forward(x, y)
+        return khd.vector_addition_forward(x, y)
 
     def backward(ctx, output_grad: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         return output_grad, output_grad