Autocast Concept: Which Ops Run in Reduced Precision, Which Stay in FP32

The problem, from first principles

02-loss-scaling fixed one specific failure mode (gradients underflowing during backward), but it didn't answer a more basic question: should EVERY operation in a model actually run in float16? 01-fp16-bf16-representable-range already showed float16 has real numerical hazards — running every single op at reduced precision would hit those hazards constantly, especially in operations that sum many values or exponentiate a wide range of numbers.

From theory to code

Implement should_run_in_fp16(op_name), which classifies a named operation as safe for float16 or requiring float32, and autocast_dtype_for(op_name), which turns that classification directly into a dtype string.

Constraints

• should_run_in_fp16 returns True for matmul-family ops (matmul, linear, conv2d, relu), False for reduction/normalization ops (softmax, log_softmax, layer_norm, batch_norm, cross_entropy, sum, mean, exp, log).
• An op name in neither list raises ValueError.
• autocast_dtype_for returns "float16" or "float32" directly, built on should_run_in_fp16.

Hints

Open one at a time. Each gives away a little more than the last.