[All] Added better error messages

[ptrendx]

Description

Added better error messages throughout the codebase in order to provide more information when the error gets triggered.

[greptile-apps]

Greptile Summary

Systematically replaced generic assertions and error messages across 26 files with detailed, actionable error messages that include actual parameter values, expected values, shapes, dtypes, and device information.

Major improvements:

• C++/CUDA error macros now specify the received dtype and list all valid options
• Python assertions converted to proper exceptions (ValueError, TypeError, RuntimeError) with context
• Error messages now include actual vs expected values for easier debugging
• Device mismatches show which device was used instead of just "needs CUDA"
• Shape mismatches include both dimensions and computed remainders
• Fixed boolean comparison in distributed.py:1873 to use if async_op: instead of identity check

Confidence Score: 5/5

• This PR is safe to merge with no risk
• Changes are purely focused on improving error messages without modifying any business logic or computational behavior. All changes convert assertions to exceptions or enhance error text with variable interpolation, making debugging easier without introducing functional changes.
• No files require special attention

Important Files Changed

Filename	Overview
transformer_engine/common/common.h	Enhanced error messages in C++ macros to specify which dtype was received and list valid options
transformer_engine/jax/flax/transformer.py	Added descriptive messages to assertion statements in Flax transformer layers
transformer_engine/pytorch/cpp_extensions/fused_attn.py	Converted assertions to exceptions with detailed shape and dtype information for fused attention
transformer_engine/pytorch/distributed.py	Replaced assertions with exceptions providing actionable error messages and actual values; fixed async_op boolean check
transformer_engine/pytorch/module/base.py	Converted assertions to detailed exceptions in module base with shape, type, and configuration information
transformer_engine/pytorch/permutation.py	Enhanced MoE permutation error messages with device, shape, and type details

_{Last reviewed commit: 221d723}

[greptile-apps]

_{26 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{26 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[ptrendx]

/te-ci

[jberchtold-nvidia]

LGTM from the TE/JAX side! Thanks! I have not reviewed the core or PyTorch files

[jberchtold-nvidia]

nit SR: If you'd like, appending "This usually occurs when TransformerEngine was not installed properly and the shared library cannot be loaded. Please see the documentation troubleshooting for assistance" but as is it's already an improvement over no message