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

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

( grid: jax.Array, rows: int, cols: int, start_row: int, start_col: int, end_row: int,

( grid: torch.Tensor, result: torch.Tensor, rows: int, cols: int, start_row: int, star

( grid: torch.Tensor, result: torch.Tensor, rows: int, cols: int, start_row: int, star

( grid: cute.Tensor, result: cute.Tensor, rows: cute.Int32, cols: cute.Int32, start_row: c

(Q: jax.Array, K: jax.Array, V: jax.Array, M: int, d: int)

(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, output: torch.Tensor, M: int, d: int)

(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, output: torch.Tensor, M: int, d: int)

( Q: cute.Tensor, K: cute.Tensor, V: cute.Tensor, output: cute.Tensor, M: cute.Int32,

( Q: jax.Array, K: jax.Array, V: jax.Array, M: int, N: int, d: int, alpha: float )

( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, output: torch.Tensor, M: int, N:

( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, output: torch.Tensor, M: int, N:

( Q: cute.Tensor, K: cute.Tensor, V: cute.Tensor, output: cute.Tensor, M: cute.Int32,

(signal: jax.Array, M: int, N: int)

(signal: torch.Tensor, spectrum: torch.Tensor, M: int, N: int)

(signal: torch.Tensor, spectrum: torch.Tensor, M: int, N: int)

(signal: cute.Tensor, spectrum: cute.Tensor, M: cute.Int32, N: cute.Int32)

(y_samples: jax.Array, a: float, b: float, n_samples: int)

(y_samples: torch.Tensor, result: torch.Tensor, a: float, b: float, n_samples: int)

(y_samples: torch.Tensor, result: torch.Tensor, a: float, b: float, n_samples: int)

( y_samples: cute.Tensor, result: cute.Tensor, a: cute.Float32, b: cute.Float32, n_samples

( A: jax.Array, B: jax.Array, M: int, N: int, K: int, scale_A: float, scale_B: flo

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

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

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

(logits: jax.Array, M: int, E: int, k: int)

( logits: torch.Tensor, topk_weights: torch.Tensor, topk_indices: torch.Tensor, M: int, E:

( logits: torch.Tensor, topk_weights: torch.Tensor, topk_indices: torch.Tensor, M: int, E:

( logits: cute.Tensor, topk_weights: cute.Tensor, topk_indices: cute.Tensor, M: cute.Int32,

( Q: jax.Array, K: jax.Array, V: jax.Array, num_q_heads: int, num_kv_heads: int, seq_l

( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, output: torch.Tensor, num_q_heads: in

( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, output: torch.Tensor, num_q_heads: in

( Q: cute.Tensor, K: cute.Tensor, V: cute.Tensor, output: cute.Tensor, num_q_heads: cute.I

(a: jax.Array, x: jax.Array, B: int, L: int)

(a: torch.Tensor, x: torch.Tensor, h: torch.Tensor, B: int, L: int)

(a: torch.Tensor, x: torch.Tensor, h: torch.Tensor, B: int, L: int)

(a: cute.Tensor, x: cute.Tensor, h: cute.Tensor, B: cute.Int32, L: cute.Int32)

(input: jax.Array, N: int)

(data: torch.Tensor, output: torch.Tensor, n: int)

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

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

(A: jax.Array, N: int)

(A: torch.Tensor, N: int, out: torch.Tensor)

(A: torch.Tensor, N: int, out: torch.Tensor)

(A: cute.Tensor, N: cute.Uint32, out: cute.Tensor)

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

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

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

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

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

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

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

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

(logits: jax.Array, p: jax.Array, seed: jax.Array, vocab_size: int)

( logits: torch.Tensor, p: torch.Tensor, seed: torch.Tensor, sampled_token: torch.Tensor,

( logits: torch.Tensor, p: torch.Tensor, seed: torch.Tensor, sampled_token: torch.Tensor,

( logits: cute.Tensor, p: cute.Tensor, seed: cute.Tensor, sampled_token: cute.Tensor, voca

( Q: jax.Array, K: jax.Array, V: jax.Array, seq_len: int, d_model: int, gamma: float,

( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, output: torch.Tensor, seq_len: int,

( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, output: torch.Tensor, seq_len: int,

( Q: cute.Tensor, K: cute.Tensor, V: cute.Tensor, output: cute.Tensor, seq_len: cute.Int32

(input: jax.Array, rows: int, cols: int)

(input: torch.Tensor, output: torch.Tensor, rows: int, cols: int)

(input: torch.Tensor, output: torch.Tensor, rows: int, cols: int)

(input: cute.Tensor, output: cute.Tensor, rows: cute.Int32, cols: cute.Int32)

( input: jax.Array, N: int, M: int, S_ROW: int, E_ROW: int, S_COL: int, E_COL: int )

( input: torch.Tensor, output: torch.Tensor, N: int, M: int, S_ROW: int, E_ROW: int,

( input: torch.Tensor, output: torch.Tensor, N: int, M: int, S_ROW: int, E_ROW: int,

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

( x: jax.Array, W_gate: jax.Array, W_up: jax.Array, W_down: jax.Array, M: int, d_model

( x: torch.Tensor, W_gate: torch.Tensor, W_up: torch.Tensor, W_down: torch.Tensor, output:

( x: torch.Tensor, W_gate: torch.Tensor, W_up: torch.Tensor, W_down: torch.Tensor, output:

( x: cute.Tensor, W_gate: cute.Tensor, W_up: cute.Tensor, W_down: cute.Tensor, output: cut

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

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

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

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

(A: jax.Array, B: jax.Array, N: int)

(a: torch.Tensor, b: torch.Tensor, result: torch.Tensor, n: int)

(A: torch.Tensor, B: torch.Tensor, result: torch.Tensor, N: int)

(A: cute.Tensor, B: cute.Tensor, result: cute.Tensor, N: cute.Int32)

(Q: jax.Array, cos: jax.Array, sin: jax.Array, M: int, D: int)

( Q: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, output: torch.Tensor, M: int, D: int )

( Q: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, output: torch.Tensor, M: int, D: int )

( Q: cute.Tensor, cos: cute.Tensor, sin: cute.Tensor, output: cute.Tensor, M: cute.Int32,

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

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

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

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

(logits: jax.Array, true_labels: jax.Array, N: int, C: int)

(logits: torch.Tensor, true_labels: torch.Tensor, loss: torch.Tensor, N: int, C: int)

(logits: torch.Tensor, true_labels: torch.Tensor, loss: torch.Tensor, N: int, C: int)

( logits: cute.Tensor, true_labels: cute.Tensor, loss: cute.Tensor, N: cute.Int32, C: cute.Int32 )

( token_ids: jax.Array, position_ids: jax.Array, token_embeddings: jax.Array, position_embeddi

( token_ids: torch.Tensor, position_ids: torch.Tensor, token_embeddings: torch.Tensor, positio

( token_ids: torch.Tensor, position_ids: torch.Tensor, token_embeddings: torch.Tensor, positio

( token_ids: cute.Tensor, position_ids: cute.Tensor, token_embeddings: cute.Tensor, position_e

(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)