code: def k_dropout_fw(
      Y,
      X,
      BIAS,
      SEEDS,
      stride,
      M,
      N,
      p: tl.constexpr,
      is_fp16: tl.constexpr,  # autotune
      ACTIVATION: tl.constexpr,
      # Meta-parameters
      BLOCK_M: tl.constexpr,
      BLOCK_N: tl.constexpr,
      SIZE_RAND_BLOCK: tl.constexpr,
      USE_BIAS: tl.constexpr,
  ):
      """
      Apply dropout on an input tensor
      Y : Output  (M, N)
      X : Input   (M, N)
      BIAS        (N,)
      SEEDS       (M,)
      p : dropout probability
      """
      # fmt: on
  
      row_id = tl.program_id(axis=0)
      rows = row_id * BLOCK_M + tl.arange(0, BLOCK_M)
  
      col_id = tl.program_id(axis=1)
      cols = col_id * BLOCK_N + tl.arange(0, BLOCK_N)
  
      # pointers starting point
      x_ptrs = X + rows[:, None] * stride + cols[None, :]
      y_ptrs = Y + rows[:, None] * stride + cols[None, :]
  
      # good to go, start the layer computations
      col_mask = cols[None, :] < N
      p_scale = 1.0 / (1.0 - p)
      if USE_BIAS:
          b_ptrs = BIAS + cols[None, :]
          bias = tl.load(b_ptrs, mask=cols[None, :] < N, other=0.0)
      else:
          bias = x_ptrs  # will not be used
  
      block_mask = (rows[:, None] < M) & col_mask
      x = tl.load(x_ptrs, mask=block_mask, other=0.0)
  
      # optionally apply a fused bias
      if USE_BIAS:
          x += bias
  
      # optional: fused activation (while the data is in shared memory)
      if ACTIVATION == 1:
          x = relu(x)
      elif ACTIVATION == 2:
          x = leaky_relu(x)
      elif ACTIVATION == 3:
          x = gelu(x)
      elif ACTIVATION == 4:
          x = squared_relu(x)
      elif ACTIVATION == 5:
          x = smelu(x)
  
      # get the random keep mask
      rand_offsets = tl.arange(0, SIZE_RAND_BLOCK)
      seed_int = tl.load(SEEDS + col_id)
      r = tl.rand(seed_int, rand_offsets)
      keep_mask = r > p
  
      # prune and normalize in one go
      keep = tl.view(keep_mask, x.shape)
      output = tl.where(keep, (x * p_scale).to(x.dtype), 0.0)
  
      tl.store(y_ptrs, output, mask=block_mask)
  
imports: import triton
  import triton.language
  from xformers.triton.k_activations import gelu, gelu_grad, leaky_relu, leaky_relu_grad, relu, relu_grad, smelu, smelu_grad, squared_relu, squared_relu_grad
label_desc: Use image node_id 1 for calling a global function with example usage: k_dropout_fw(Y, X, BIAS, SEEDS, stride, M, N, p, is_fp16, ACTIVATION, BLOCK_M, BLOCK_N, SIZE_RAND_BLOCK, USE_BIAS) without return types