Ensure tests also work in test shells

tests/test_rotary.py

@@ -2,7 +2,6 @@ import pytest
 import torch
 
 from tests.utils import infer_device, supports_bfloat16
-from kernels import get_local_kernel
 from pathlib import Path
 
 # import rotary
@@ -10,8 +9,12 @@ from pathlib import Path
 # set_seed(42)
 
 # Set the local repo path, relative path
-repo_path = Path(__file__).parent.parent
-rotary = get_local_kernel(repo_path=repo_path, package_name="rotary")
+try:
+    import rotary
+except ImportError:
+    from kernels import get_local_kernel
+    repo_path = Path(__file__).parent.parent
+    rotary = get_local_kernel(repo_path=repo_path, package_name="rotary")
 
 def apply_rotary_torch(x1: torch.Tensor, x2: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, conj: bool = False):
     assert x1.shape == x2.shape, "x1 and x2 must have the same shape"