[Nebulixeay]

Problem

With autocast enabled, Linear and Matmul can return lower-precision tensors. The recorded MLU BF16 GPT-2 path encountered downstream gradient dtype mismatches. This change represents the lower-precision operator outputs as FP32 for the downstream path.

Change

After SetupContext and before saved-variable wiring, convert outputs of the configured lower-precision autocast operators to FP32 when their dtype matches the active autocast dtype. Transfer non-differentiable output markings to the replacement tensor. Disabled autocast, unrelated operators, null outputs, and outputs with other dtypes keep their existing tensors.

The conversion preserves the rounded lower-precision values; it does not recover FP32 precision. Only infini_train/src/autograd/function.cc changes, including its direct includes and formatting. The framework API and operator list are unchanged.

Validation

• The functional change was built and exercised in downstream MLU BF16 single-step and ten-step diagnostic runs.
• The follow-up includes/comment/formatting cleanup preserves executable code after excluding the added includes, comments, and whitespace.
• clang-format 16.0.6 single-file check and git diff --check pass.
• GitHub Format Check, IWYU, and clang-tidy results must be evaluated for the latest commit.
• No new device training run was performed for the formatting cleanup. The recorded runs are diagnostic evidence, not formal cross-framework accuracy acceptance or cross-provider validation.