Change logging to logger

Author: 0x45f

src/flag_gems/ops/baddbmm.py

@@ -16,6 +16,8 @@
 
 from .mul import mul
 
+logger = logging.getLogger(__name__)
+
 
 @libentry()
 @libtuner(
@@ -135,7 +137,7 @@ def baddbmm_kernel(
 class BaddbmmFunction(torch.autograd.Function):
     @staticmethod
     def forward(ctx, bias, A, B, beta, alpha):
-        logging.debug("GEMS BADDBMM FORWARD")
+        logger.debug("GEMS BADDBMM FORWARD")
 
         ctx.save_for_backward(A, B, bias)
         ctx.alpha = alpha
@@ -176,7 +178,7 @@ def forward(ctx, bias, A, B, beta, alpha):
 
     @staticmethod
     def backward(ctx, grad_output):
-        logging.debug("GEMS BADDBMM BACKWARD")
+        logger.debug("GEMS BADDBMM BACKWARD")
         A, B, bias = ctx.saved_tensors
 
         grad_A = None

src/flag_gems/ops/fill.py

@@ -60,7 +60,7 @@ def fill_tensor_(self, value):
 
 
 def fill_scalar_(self, value):
-    logging.debug("GEMS FILL_SCALAR_")
+    logger.debug("GEMS FILL_SCALAR_")
     with torch_device_fn.device(self.device):
         fill_scalar_func(self, value, out0=self)
     return self

src/flag_gems/runtime/backend/_ascend/fused/rotary_embedding.py

@@ -188,7 +188,7 @@ def apply_rotary_pos_emb(
         q_embed: (*, q_heads, head_dim)
         k_embed: (*, k_heads, head_dim)
     """
-    logging.debug("GEMS_ASCEND ROTARY POS EMBEDDING")
+    logger.debug("GEMS_ASCEND ROTARY POS EMBEDDING")
     assert (
         k.shape[-1] == q.shape[-1]
     ), f"q and k must have the same last dimension, got {q.shape} and {k.shape}"

src/flag_gems/runtime/backend/_ascend/ops/groupnorm.py

@@ -246,7 +246,6 @@ def group_norm_kernel(
 
 
 def group_norm(input, weight, bias, N, C, HxW, group, eps=1e-05):
-    logging.debug("ASCEND GEMS GROUPNORM FORWARD")
     logger.debug("ASCEND GEMS GROUPNORM FORWARD")
     group_size = triton.cdiv(C, group)
     input = input.contiguous()
@@ -280,7 +279,6 @@ def group_norm(input, weight, bias, N, C, HxW, group, eps=1e-05):
 def group_norm_backward(
     grad_out, input, mean, rstd, weight, N, C, HxW, group, output_mask
 ):
-    logging.debug("ASCEND GEMS GROUPNORM BACKWARD")
     logger.debug("ASCEND GEMS GROUPNORM BACKWARD")
     grad_out = grad_out.contiguous()
     input = input.contiguous()

src/flag_gems/runtime/backend/_ascend/ops/hstack.py

@@ -20,7 +20,7 @@ def copy_func(x):
 def hstack(
     tensors: Union[Tuple[torch.Tensor, ...], List[torch.Tensor]]
 ) -> torch.Tensor:
-    logging.debug("GEMS_ASCEND HSTACK")
+    logger.debug("GEMS_ASCEND HSTACK")
 
     if len(tensors) == 0:
         raise RuntimeError("hstack expected a non-empty TensorList")
@@ -70,4 +70,4 @@ def hstack(
         out_view = StridedBuffer(out0, a.shape, out0.stride(), offset=out0_offset)
         copy_func.instantiate(a.ndim)(in_view, out0=out_view)
 
-    return out0
+    return out0