Custom Loss Function AI Prompt

Implement a custom loss function for this problem with full PyTorch autograd compatibility.

Problem: {{problem_description}}
Loss requirements: {{loss_requirements}}

1. Implementation requirements:
   - Subclass torch.nn.Module or implement as a function
   - Fully differentiable — use only PyTorch tensor operations, no NumPy inside forward()
   - Support batched inputs of arbitrary batch size
   - Handle edge cases: empty batches, all-same-class batches, NaN/Inf inputs

2. For composite losses (combining multiple terms):
   - Implement each term as a separate method for testability
   - Use learnable or fixed weighting between terms
   - Log each term's contribution separately during training

3. Numerical stability:
   - Use log-sum-exp trick for log probabilities
   - Apply clipping to prevent log(0) or division by zero
   - Test with fp16 — ensure no overflow with half precision

4. Testing the loss:
   - Unit test: verify loss = 0 for perfect predictions
   - Unit test: verify loss > 0 for wrong predictions
   - Gradient check: torch.autograd.gradcheck to verify analytical gradients match numerical approximation
   - Verify loss is lower for better predictions (sanity check)

5. Reduction modes: support 'mean', 'sum', and 'none' as in standard PyTorch losses

Return: loss implementation, unit tests, gradient check, and integration example in a training loop.