# Generated from the displayed GPU Lab example; no GPU timings are measured.
# Run: python example.py --device cpu  (default; no CPU/GPU transfers)
# Or:  python example.py --device cuda (compatible NVIDIA GPU and PyTorch required)
# Explicit CPU-only calculations remain on CPU in either mode.
import argparse
import torch

_lab_parser = argparse.ArgumentParser(description="Run this GPU Lab PyTorch example.")
_lab_parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
_lab_device = _lab_parser.parse_args().device
if _lab_device == "cuda" and not torch.cuda.is_available():
    _lab_parser.error("CUDA mode needs an available NVIDIA GPU and compatible PyTorch installation.")
if _lab_device == "cpu":
    print("CPU mode checks values and layouts; it does not perform CPU-GPU transfers.")

# The displayed example, with its CUDA device requests selected by --device.
import torch
import torch.nn.functional as F

# Three cached tokens, each with four components.
k_cache = torch.arange(12, dtype=torch.float32, device=_lab_device).reshape(3, 4) / 10
v_cache = torch.arange(12, dtype=torch.float32, device=_lab_device).reshape(3, 4)
q_new = torch.tensor([[1., 0., -1., 0.]], dtype=torch.float32, device=_lab_device)
k_new = torch.tensor([[0.2, 0.3, 0.4, 0.5]], dtype=torch.float32, device=_lab_device)
v_new = torch.tensor([[12., 13., 14., 15.]], dtype=torch.float32, device=_lab_device)
k_cache = torch.cat((k_cache, k_new), dim=-2)
v_cache = torch.cat((v_cache, v_new), dim=-2)
out = F.scaled_dot_product_attention(
    q_new, k_cache, v_cache, dropout_p=0.0, is_causal=False)
print(k_cache.shape, out)