def fused_fwd_bwd(
            input_chunk,
            selected_token_ids_chunk,
            attention_mask_chunk,
            advantages_chunk,
            ref_per_token_logps_chunk,
            old_per_token_logps_chunk,
            ref_input_chunk,
        ):
            """Fused forward and backward for a chunk."""
            argnums = (0, 1, 5) if bias is not None else (0, 1)
            return torch.func.grad_and_value(compute_loss, argnums=argnums, has_aux=True)(
                input_chunk,  # arg 0
                weight,  # arg 1
                selected_token_ids_chunk,  # arg 2
                attention_mask_chunk,  # arg 3
                advantages_chunk,  # arg 4
                bias,  # arg 5
                ref_per_token_logps_chunk=ref_per_token_logps_chunk,  # arg 6
                old_per_token_logps_chunk=old_per_token_logps_chunk,  # arg 7
                ref_input_chunk=ref_input_chunk,  # arg 8
            )
Fused forward and backward for a chunk.
src/liger_kernel/chunked_loss/fused_linear_ppo.py