(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, BATCH: int, M: int, N: int, K: int)

( A: cute.Tensor, B: cute.Tensor, C: cute.Tensor, BATCH: cute.Int32, M: cute.Int32, N:

(X: jax.Array, S: jax.Array, M: int, N: int, TILE_SIZE: int)

(X: torch.Tensor, S: torch.Tensor, Y: torch.Tensor, M: int, N: int, TILE_SIZE: int)

(X: torch.Tensor, S: torch.Tensor, Y: torch.Tensor, M: int, N: int, TILE_SIZE: int)

( X: cute.Tensor, S: cute.Tensor, Y: cute.Tensor, M: cute.Int32, N: cute.Int32, TILE_S

(A: jax.Array, B: jax.Array, M: int, N: int, K: int, nnz: int)

(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, M: int, N: int, K: int, nnz: int)

(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, M: int, N: int, K: int, nnz: int)

( A: cute.Tensor, B: cute.Tensor, C: cute.Tensor, M: cute.Int32, N: cute.Int32, K: cut

( u: jax.Array, delta: jax.Array, A: jax.Array, B: jax.Array, C: jax.Array, skip: jax.

( u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,

( u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,

( u: cute.Tensor, delta: cute.Tensor, A: cute.Tensor, B: cute.Tensor, C: cute.Tensor,

( input: jax.Array, kernel: jax.Array, input_rows: int, input_cols: int, kernel_rows: int,

( input: torch.Tensor, kernel: torch.Tensor, output: torch.Tensor, input_rows: int, input_

( input: torch.Tensor, kernel: torch.Tensor, output: torch.Tensor, input_rows: int, input_

( input: cute.Tensor, kernel: cute.Tensor, output: cute.Tensor, input_rows: cute.Int32, in

( input: jax.Array, kernel: jax.Array, input_rows: int, input_cols: int, kernel_rows: int,

( input: torch.Tensor, kernel: torch.Tensor, output: torch.Tensor, input_rows: int, input_

( input: torch.Tensor, kernel: torch.Tensor, output: torch.Tensor, input_rows: int, input_

( input: cute.Tensor, kernel: cute.Tensor, output: cute.Tensor, input_rows: cute.Int32, in

(input: jax.Array, N: int, window_size: int)

(input: torch.Tensor, output: torch.Tensor, N: int, window_size: int)

(input: torch.Tensor, output: torch.Tensor, N: int, window_size: int)

(input: cute.Tensor, output: cute.Tensor, N: cute.Int32, window_size: cute.Int32)

( x: jax.Array, w_q: jax.Array, scales: jax.Array, M: int, N: int, K: int, group_s

( x: torch.Tensor, w_q: torch.Tensor, scales: torch.Tensor, y: torch.Tensor, M: int, N

( x: torch.Tensor, w_q: torch.Tensor, scales: torch.Tensor, y: torch.Tensor, M: int, N

( x: cute.Tensor, w_q: cute.Tensor, scales: cute.Tensor, y: cute.Tensor, M: cute.Int32,

( x: jax.Array, W: jax.Array, A: jax.Array, B: jax.Array, batch: int, d_in: int, d

( x: torch.Tensor, W: torch.Tensor, A: torch.Tensor, B: torch.Tensor, output: torch.Tensor

( x: torch.Tensor, W: torch.Tensor, A: torch.Tensor, B: torch.Tensor, output: torch.Tensor

( x: cute.Tensor, W: cute.Tensor, A: cute.Tensor, B: cute.Tensor, output: cute.Tensor,

(input: jax.Array, gamma: jax.Array, beta: jax.Array, N: int, eps: float)

(input: torch.Tensor, gamma: float, beta: float, output: torch.Tensor, N: int, eps: float)

( input: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, output: torch.Tensor, N: i

( input: cute.Tensor, gamma: cute.Float32, beta: cute.Float32, output: cute.Tensor, N: cut

( Q: jax.Array, K_int8: jax.Array, V_int8: jax.Array, k_scale: jax.Array, v_scale: jax.Arr

( Q: torch.Tensor, K_int8: torch.Tensor, V_int8: torch.Tensor, k_scale: torch.Tensor, v_sc

( Q: torch.Tensor, K_int8: torch.Tensor, V_int8: torch.Tensor, k_scale: torch.Tensor, v_sc

( Q: cute.Tensor, K_int8: cute.Tensor, V_int8: cute.Tensor, k_scale: cute.Tensor, v_scale:

( draft_tokens: jax.Array, draft_probs: jax.Array, target_probs: jax.Array, uniform_samples: j

( draft_tokens: torch.Tensor, draft_probs: torch.Tensor, target_probs: torch.Tensor, uniform_s

( draft_tokens: torch.Tensor, draft_probs: torch.Tensor, target_probs: torch.Tensor, uniform_s

( draft_tokens: cute.Tensor, draft_probs: cute.Tensor, target_probs: cute.Tensor, uniform_samp

( input: jax.Array, N: int, C: int, H: int, W: int, kernel_size: int, stride: int, padding: int )

(input, output, N, C, H, W, kernel_size, stride, padding)

(input, output, N, C, H, W, kernel_size, stride, padding)

( input: cute.Tensor, output: cute.Tensor, N: cute.Int32, C: cute.Int32, H: cute.Int32,

( X: jax.Array, gamma: jax.Array, beta: jax.Array, N: int, C: int, H: int, W: int,

( X: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, Y: torch.Tensor, N: int, C

( X: torch.Tensor, gamma: torch.Tensor, beta: torch.Tensor, Y: torch.Tensor, N: int, C

( X: cute.Tensor, gamma: cute.Tensor, beta: cute.Tensor, Y: cute.Tensor, N: cute.Int32,

(input: jax.Array, N: int, S: int, E: int)

(input: torch.Tensor, output: torch.Tensor, N: int, S: int, E: int)

(input: torch.Tensor, output: torch.Tensor, N: int, S: int, E: int)

(input: cute.Tensor, output: cute.Tensor, N: cute.Int32, S: cute.Int32, E: cute.Int32)

( x: jax.Array, weight: jax.Array, bias: jax.Array, B: int, L: int, D: int, K: int )

( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, output: torch.Tensor, B: int,

( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, output: torch.Tensor, B: int,

( x: cute.Tensor, weight: cute.Tensor, bias: cute.Tensor, output: cute.Tensor, B: cute.Int

(prompts: jax.Array, weights: jax.Array, batch_size: int)

(prompts: torch.Tensor, output: torch.Tensor, weights: torch.Tensor, batch_size: int)

(prompts: torch.Tensor, output: torch.Tensor, weights: torch.Tensor, batch_size: int)

( prompts: cute.Tensor, output: cute.Tensor, weights: cute.Tensor, batch_size: cute.Int32, )

(input: jax.Array, N: int, num_bins: int)

(input: torch.Tensor, histogram: torch.Tensor, N: int, num_bins: int)

(input: torch.Tensor, histogram: torch.Tensor, N: int, num_bins: int)

(input: cute.Tensor, histogram: cute.Tensor, N: cute.Int32, num_bins: cute.Int32)

(input: jax.Array, N: int)

(input: torch.Tensor, output: torch.Tensor, N: int)

(input: torch.Tensor, output: torch.Tensor, N: int)

(input: cute.Tensor, output: cute.Tensor, N: cute.Int32)

( input: jax.Array, kernel: jax.Array, input_depth: int, input_rows: int, input_cols: int,

( input: torch.Tensor, kernel: torch.Tensor, output: torch.Tensor, input_depth: int, input

( input: torch.Tensor, kernel: torch.Tensor, output: torch.Tensor, input_depth: int, input

( input: cute.Tensor, kernel: cute.Tensor, output: cute.Tensor, input_depth: cute.Int32, i

(input: jax.Array, N: int, M: int, K: int)

(input: torch.Tensor, output: torch.Tensor, N: int, M: int, K: int)

(input: torch.Tensor, output: torch.Tensor, N: int, M: int, K: int)

(input: cute.Tensor, output: cute.Tensor, N: cute.Int32, M: cute.Int32, K: cute.Int32)

(input: jax.Array, N: int, M: int, K: int, P: int)

(input: torch.Tensor, output: torch.Tensor, N: int, M: int, K: int, P: int)

(input: torch.Tensor, output: torch.Tensor, N: int, M: int, K: int, P: int)

( input: cute.Tensor, output: cute.Tensor, N: cute.Int32, M: cute.Int32, K: cute.Int32,

( input: jax.Array, N: int, M: int, K: int, S_DEP: int, E_DEP: int, S_ROW: int,

( input: torch.Tensor, output: torch.Tensor, N: int, M: int, K: int, S_DEP: int, E

( input: torch.Tensor, output: torch.Tensor, N: int, M: int, K: int, S_DEP: int, E

( input: cute.Tensor, output: cute.Tensor, N: cute.Int32, M: cute.Int32, K: cute.Int32,

(input: jax.Array, N: int)

(input: torch.Tensor, output: torch.Tensor, N: int)

(input: torch.Tensor, output: torch.Tensor, N: int)

(input: cute.Tensor, output: cute.Tensor, N: cute.Int32)

(predictions: jax.Array, targets: jax.Array, N: int)

(predictions: torch.Tensor, targets: torch.Tensor, mse: torch.Tensor, N: int)

(predictions: torch.Tensor, targets: torch.Tensor, mse: torch.Tensor, N: int)

(predictions: cute.Tensor, targets: cute.Tensor, mse: cute.Tensor, N: cute.Int32)

(A: jax.Array, B: jax.Array, BATCH: int, M: int, N: int, K: int)

(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor, BATCH: int, M: int, N: int, K: int)