Avoid relative imports (#739)

* avoid relative imports in python files that contain triton jit functions to allow importing those files as top-level files

Author: iclementine

src/flag_gems/fused/__init__.py

@@ -1,14 +1,14 @@
-from .concat_and_cache_mla import concat_and_cache_mla
-from .cross_entropy_loss import cross_entropy_loss
-from .fused_add_rms_norm import fused_add_rms_norm
-from .gelu_and_mul import gelu_and_mul
-from .instance_norm import instance_norm
-from .outer import outer
-from .reshape_and_cache import reshape_and_cache
-from .rotary_embedding import apply_rotary_pos_emb
-from .silu_and_mul import silu_and_mul
-from .skip_layernorm import skip_layer_norm
-from .weight_norm import weight_norm
+from flag_gems.fused.concat_and_cache_mla import concat_and_cache_mla
+from flag_gems.fused.cross_entropy_loss import cross_entropy_loss
+from flag_gems.fused.fused_add_rms_norm import fused_add_rms_norm
+from flag_gems.fused.gelu_and_mul import gelu_and_mul
+from flag_gems.fused.instance_norm import instance_norm
+from flag_gems.fused.outer import outer
+from flag_gems.fused.reshape_and_cache import reshape_and_cache
+from flag_gems.fused.rotary_embedding import apply_rotary_pos_emb
+from flag_gems.fused.silu_and_mul import silu_and_mul
+from flag_gems.fused.skip_layernorm import skip_layer_norm
+from flag_gems.fused.weight_norm import weight_norm
 
 __all__ = [
     "apply_rotary_pos_emb",

src/flag_gems/fused/concat_and_cache_mla.py

@@ -4,7 +4,7 @@
 import triton
 import triton.language as tl
 
-from ..utils import libentry
+from flag_gems.utils import libentry
 
 logger = logging.getLogger(__name__)
 

src/flag_gems/fused/cross_entropy_loss.py

@@ -5,10 +5,10 @@
 import triton.language as tl
 from torch.nn import _reduction as _Reduction
 
-from .. import runtime
-from ..runtime import torch_device_fn
-from ..utils import libentry
-from ..utils import triton_lang_extension as tle
+from flag_gems import runtime
+from flag_gems.runtime import torch_device_fn
+from flag_gems.utils import libentry
+from flag_gems.utils import triton_lang_extension as tle
 
 logger = logging.getLogger(__name__)
 

src/flag_gems/fused/fused_add_rms_norm.py

@@ -4,9 +4,9 @@
 import triton
 import triton.language as tl
 
-from ..runtime import torch_device_fn
-from ..utils import libentry
-from ..utils import triton_lang_extension as tle
+from flag_gems.runtime import torch_device_fn
+from flag_gems.utils import libentry
+from flag_gems.utils import triton_lang_extension as tle
 
 logger = logging.getLogger(__name__)
 

src/flag_gems/fused/gelu_and_mul.py

@@ -4,7 +4,7 @@
 import triton
 import triton.language as tl
 
-from ..utils import pointwise_dynamic, tl_extra_shim
+from flag_gems.utils import pointwise_dynamic, tl_extra_shim
 
 erf = tl_extra_shim.erf
 pow = tl_extra_shim.pow

src/flag_gems/fused/instance_norm.py

@@ -6,10 +6,10 @@
 import triton
 import triton.language as tl
 
-from .. import runtime
-from ..runtime import torch_device_fn
-from ..utils import libentry
-from ..utils.type_utils import get_accumulator_dtype
+from flag_gems import runtime
+from flag_gems.runtime import torch_device_fn
+from flag_gems.utils import libentry
+from flag_gems.utils.type_utils import get_accumulator_dtype
 
 logger = logging.getLogger(__name__)
 Tensor = torch.Tensor

src/flag_gems/fused/outer.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from ..ops import mul, mv
+from flag_gems.ops import mul, mv
 
 logger = logging.getLogger(__name__)
 

src/flag_gems/fused/reshape_and_cache.py

@@ -3,8 +3,8 @@
 import triton
 import triton.language as tl
 
-from ..runtime import torch_device_fn
-from ..utils import libentry
+from flag_gems.runtime import torch_device_fn
+from flag_gems.utils import libentry
 
 logger = logging.getLogger(__name__)
 

src/flag_gems/fused/rotary_embedding.py

@@ -5,9 +5,9 @@
 import triton
 import triton.language as tl
 
-from ..runtime import torch_device_fn
-from ..utils import libentry
-from ..utils import triton_lang_extension as tle
+from flag_gems.runtime import torch_device_fn
+from flag_gems.utils import libentry
+from flag_gems.utils import triton_lang_extension as tle
 
 logger = logging.getLogger(__name__)
 

src/flag_gems/fused/silu_and_mul.py

@@ -4,7 +4,7 @@
 import triton
 import triton.language as tl
 
-from ..utils import pointwise_dynamic
+from flag_gems.utils import pointwise_dynamic
 
 logger = logging.getLogger(__name__)