print(
    "COMPILE_OPTIONS_JSON="
    + json.dumps(cross_entropy_loss.torch_compile_options, sort_keys=True)
)

torch._dynamo.reset()
torch._dynamo.utils.counters.clear()
torch._dynamo.utils.guard_failures.clear()
cross_entropy_loss._FUSED_CE_COMPILE_SUPPORTED = None
cross_entropy_loss._FUSED_CE_COMPILE_FASTPATH_PROVEN.clear()

# Two equal chunks exercise one compiled accumulate_chunk region twice. Each
# chunk has one valid and one ignored row. All inputs are finite, but the
# ignored rows overflow in linear(), which makes the sanitizer observable.
hidden_states = torch.tensor(
    [[[2.0e-20, -1.0e-20], [1.0e20, -1.0e20],
      [1.5e-20, -0.5e-20], [1.0e20, -1.0e20]]],
    dtype=torch.float32,
    device="cuda",
    requires_grad=True,
)
lm_head_weight = torch.tensor(
    [[1.0e20, 0.0], [0.0, 1.0e20], [-1.0e20, 1.0e20]],
    dtype=torch.float32,
    device="cuda",
    requires_grad=True,
)
lm_head_bias = torch.tensor(
    [0.1, -0.2, 0.3],
    dtype=torch.float32,
    device="cuda",
    requires_grad=True,
)
labels = torch.tensor([[1, -100, 1, -100]], device="cuda")

with torch.no_grad():
    raw_logits = torch.nn.functional.linear(
        hidden_states,
        lm_head_weight,
        lm_head_bias,
    )
print("RAW_LOGITS=" + repr(raw_logits.detach().cpu().tolist()))

loss = cross_entropy_loss.unsloth_fused_dft_loss(
    trainer=None,
    hidden_states=hidden_states,
    lm_head_weight=lm_head_weight,
    lm_head_bias=lm_head_bias,
    labels=labels,
    shift_labels=False,
    torch_compile=True,
    n_chunks=2,
)
torch.cuda.synchronize()
loss.backward()
torch.cuda.synchronize()

guard_failures = {
    str(code): [str(failure) for failure in failures]
    for code, failures in torch._dynamo.utils.guard_failures.items()
    if failures
}
counters = json_counters()
graph_break_count = sum(counters.get("graph_break", {}).values())
recompilation_count = sum(len(failures) for failures in guard_failures.values())

result = {
    "loss": float(loss.detach().cpu()),
    "compile_supported": cross_entropy_loss._FUSED_CE_COMPILE_SUPPORTED,
    "fastpath_proven": (
        cross_entropy_loss.compute_fused_dft_loss
        in cross_entropy_loss._FUSED_CE_COMPILE_FASTPATH_PROVEN
    ),
    "hidden_grad_finite": bool(torch.isfinite(hidden_states.grad).all()),
    "weight_grad_finite": bool(torch.isfinite(lm_head_weight.grad).all()),
    "bias_grad_finite": bool(torch.isfinite(lm_head_bias.grad).all()),
    "ignored_hidden_grad_zero": bool(
        torch.count_nonzero(hidden_states.grad[0, [1, 3]]) == 0
    ),
    "graph_break_count": graph_break_count,
    "recompilation_count_from_guard_failures": recompilation_count,
    "guard_failures": guard_failures,
    "counters": counters,
}
print("RESULT_JSON=" + json.dumps(result, sort_keys=True))

COMPILE_OPTIONS_JSON={"benchmark_combo_kernel": true, "combo_kernel_foreach_dynamic_shapes": true, "combo_kernels": false, "compile_threads": 16, "coordinate_descent_tuning": false, "cuda.compile_opt_level": "-O2", "cuda.enable_cuda_lto": true, "dce": true, "debug": true, "disable_progress": true, "epilogue_fusion": true, "group_fusion": true, "max_autotune": false, "memory_planning": false, "shape_padding": true, "trace.enabled": true, "trace.graph_diagram": true, "triton.autotune_at_compile_time": false, "triton.cooperative_reductions": false, "triton.cudagraphs": false, "triton.enable_persistent_tma_matmul": true, "triton.multi_kernel": 0, "triton.use_block_ptr": false, "verbose_progress": false}

RAW_LOGITS=[[[2.0999999046325684, -1.2000000476837158, -2.700000047683716], [inf, -inf, -inf], [1.600000023841858, -0.699999988079071, -1.7000000476837158], [inf, -inf, -inf]]]

RESULT_JSON={"bias_grad_finite": true, "compile_supported": true, "counters": {"aot_autograd": {"autograd_cache_miss": 1, "autograd_cache_saved": 1, "ok": 1, "total": 1}, "aten_mm_info": {"aten.mm_s19_s19_s55": 1, "aten.mm_s19_s55_s19": 1, "aten.mm_s55_s19_s19": 1}, "inductor": {"async_compile_cache_hit": 6, "async_compile_cache_miss": 12, "benchmarking.InductorBenchmarker.benchmark": 5, "benchmarking.InductorBenchmarker.benchmark_gpu": 5, "extern_calls": 3, "fxgraph_cache_miss": 1, "pattern_matcher_count": 7, "pattern_matcher_nodes": 12, "triton_bundler_save_kernel": 72, "triton_bundler_save_static_autotuner": 1}, "stats": {"calls_captured": 59, "unique_graphs": 1}}, "fastpath_proven": true, "graph_break_count": 0, "guard_failures": {}, "hidden_grad_finite": true, "ignored_hidden_grad_zero": true, "loss": 0.1660669595003128, "recompilation_count_from_guard_failures": 0, "weight_grad_finite": true}