Revert "update tests"

This reverts commit 11343cd.
DeepLink-org · Nov 5, 2024 · da385c3 · da385c3
1 parent 11343cd
commit da385c3
Show file tree

Hide file tree

Showing 8 changed files with 15 additions and 15 deletions.
diff --git a/tests/easyllm/test_rms_norm_dipu.py b/tests/easyllm/test_rms_norm_dipu.py
@@ -2,8 +2,8 @@
 
 import torch
 from tests.core import calculate_fwd_and_bwd, allclose
-from deeplink_ext.ops.rms_norm import rms_norm
-from deeplink_ext.ops.rms_norm.easyllm_rms_norm_fallback import rms_norm_torch
+from deeplink_ext.easyllm_ops.rms_norm import rms_norm
+from deeplink_ext.easyllm_ops.rms_norm_fallback import rms_norm_torch
 
 
 def test_rms_norm():

diff --git a/tests/internevo/test_flash_attention.py b/tests/internevo/test_flash_attention.py
@@ -3,12 +3,12 @@
 import torch
 from tests.core import copy_to_cpu, allclose, calculate_fwd_and_bwd
 
-from deeplink_ext.ops.flash_attention.internevo_flash_attention_fallback  import (
+from deeplink_ext.internevo_ops.flash_attention_fallback import (
     flash_attn_qkvpacked_func_torch,
     flash_attn_kvpacked_func_torch,
     flash_attn_func_torch,
 )
-from deeplink_ext.ops.flash_attention.internevo_flash_attention import (
+from deeplink_ext.internevo_ops.flash_attention import (
     flash_attn_qkvpacked_func,
     flash_attn_kvpacked_func,
     flash_attn_func,

diff --git a/tests/internevo/test_rotary_embedding.py b/tests/internevo/test_rotary_embedding.py
@@ -2,8 +2,8 @@
 
 import torch
 from tests.core import call_autograd_func, allclose
-from deeplink_ext.ops.rotary_embedding.internevo_rotary_embedding import ApplyRotaryEmb
-from deeplink_ext.ops.rotary_embedding.internevo_rotary_embedding_fallback import ApplyRotaryEmbTorch
+from deeplink_ext.internevo_ops.rotary_embedding import ApplyRotaryEmb
+from deeplink_ext.internevo_ops.rotary_embedding_fallback import ApplyRotaryEmbTorch
 
 
 def test_ApplyRotaryEmb():

diff --git a/tests/internevo/test_varlen_flash_attention.py b/tests/internevo/test_varlen_flash_attention.py
@@ -3,12 +3,12 @@
 import torch
 from tests.core import allclose, calculate_fwd_and_bwd, copy_to_cpu
 
-from deeplink_ext.ops.flash_attention.internevo_flash_attention_fallback import (
+from deeplink_ext.internevo_ops.flash_attention_fallback import (
     flash_attn_varlen_qkvpacked_func_torch,
     flash_attn_varlen_kvpacked_func_torch,
     flash_attn_varlen_func_torch,
 )
-from deeplink_ext.ops.flash_attention.internevo_flash_attention import (
+from deeplink_ext.internevo_ops.flash_attention import (
     flash_attn_varlen_qkvpacked_func,
     flash_attn_varlen_kvpacked_func,
     flash_attn_varlen_func,

diff --git a/tests/interntrain/test_adamw_dipu.py b/tests/interntrain/test_adamw_dipu.py
@@ -3,7 +3,7 @@
 import copy
 import torch
 from torch import nn
-from deeplink_ext.ops.adamw import AdamW
+from deeplink_ext.interntrain_ops.adamw import AdamW
 
 
 def test_AdamW():

diff --git a/tests/interntrain/test_flash_attention.py b/tests/interntrain/test_flash_attention.py
@@ -3,11 +3,11 @@
 import torch
 from tests.core import copy_to_cpu, allclose, call_module
 
-from deeplink_ext.ops.flash_attention.interntrain_flash_attention import (
+from deeplink_ext.interntrain_ops.flash_attention import (
     FlashSelfAttention,
     FlashCrossAttention,
 )
-from deeplink_ext.ops.flash_attention.interntrain_flash_attention_fallback import (
+from deeplink_ext.interntrain_ops.flash_attention_fallback import (
     SelfAttention,
     CrossAttention,
 )

diff --git a/tests/interntrain/test_rms_norm.py b/tests/interntrain/test_rms_norm.py
@@ -2,8 +2,8 @@
 
 import torch
 from tests.core import call_module, allclose
-from deeplink_ext.ops.rms_norm.internevo_rms_norm import MixedFusedRMSNorm
-from deeplink_ext.ops.rms_norm.internevo_rms_norm_fallback import MixedRMSNormTorch
+from deeplink_ext.interntrain_ops.rms_norm import MixedFusedRMSNorm
+from deeplink_ext.interntrain_ops.rms_norm_fallback import MixedRMSNormTorch
 
 
 def test_MixedFusedRMSNorm():

diff --git a/tests/interntrain/test_rotary_embedding.py b/tests/interntrain/test_rotary_embedding.py
@@ -3,11 +3,11 @@
 import torch
 from tests.core import call_autograd_func, allclose
 
-from deeplink_ext.ops.rotary_embedding.interntrain_rotary_embedding import (
+from deeplink_ext.interntrain_ops.rotary_embedding import (
     ApplyRotaryEmb,
     ApplyRotaryEmbQKV_,
 )
-from deeplink_ext.ops.rotary_embedding.interntrain_rotary_embedding_fallback import (
+from deeplink_ext.interntrain_ops.rotary_embedding_fallback import (
     ApplyRotaryEmbTorch,
     ApplyRotaryEmbQKV_Torch,
 )