feat(casting): support ignored_ops in cast_fp32_to_fp16 and cast_to_16_bit_precision (#7)

[rohith500]

Summary

This PR adds an ignored_ops parameter to cast_fp32_to_fp16 and cast_to_16_bit_precision, allowing users to specify operations that should be preserved in torch.float32 during 16-bit casting. This resolves #7 by providing a deterministic mechanism to prevent intermediate activation-level overflow in operations such as exp, logsumexp, softplus, or mish.

Motivation & Background

As discussed in #7, IEEE 754 half-precision float (FP16) has a maximum finite limit of 65,504. Neural network activation functions involving exponentials (such as $\text{Softplus}(x) = \ln(1 + \exp(x))$, $\text{LogSumExp}$, or $\text{Mish}$) exceed this limit whenever intermediate activations $x > 11.0$ (e.g. $\exp(15.0) \approx 3.27 \times 10^6$), causing the operation to evaluate to $+\infty$. On hardware backends like Apple Silicon / ANE, this overflow causes numerical collapse downstream into 0.0 or NaN.

Because cast_fp32_to_fp16 performs graph transformations prior to runtime without inspecting dynamic activation magnitudes, it eagerly converts these compute ops to torch.float16. Following maintainer guidance from @crowbat in #7 (Direction 2: "allowing the user to specify specific ops or op types to ignore being casted"), this PR enables users to selectively skip casting for sensitive operations.

Changes

1. Selective Op Identification (is_ignored_op):

   • Added is_ignored_op in coreai_opt._utils.casting_utils.
   • Supports 5 flexible identifier formats:
     • torch._ops.OpOverload (e.g. torch.ops.aten.exp.default)
     • torch._ops.OpOverloadPacket (e.g. torch.ops.aten.exp)
     • Standard PyTorch callables (e.g. torch.exp)
     • String op names (e.g. "exp", "aten.exp", "aten::exp")
     • String node names (e.g. "exp_1", targeting specific graph nodes)

2. Casting Engine Integration (_FP16Casting):

   • In _FP16Casting.iterate_nodes_and_insert_casts, checks is_ignored_op(node, self._ignored_ops) before creation and general compute ops.
   • Routes ignored ops directly to handle_overflow_op(node), which preserves the op in torch.float32, inserts FP16 -> FP32 upcasts on any FP16 inputs, and inserts FP32 -> FP16 downcasts on outputs for downstream consumers.
   • Redundant intermediate casts between consecutive ignored ops (e.g. exp $\to$ log1p) are automatically collapsed by the existing _cleanup_casts engine into a direct FP32 chain with zero overhead.

3. Public API Updates:

   • Added ignored_ops: Sequence[Callable | str] | None = None to cast_fp32_to_fp16 and cast_to_16_bit_precision.
   • Fully backward-compatible; defaults to None.

4. Testing & Documentation:

   • Added TestIsIgnoredOp unit tests in tests/casting/test_casting_utils.py covering all matching styles and edge cases.
   • Added TestSelectiveOpSkipping end-to-end tests in tests/casting/test_casting.py verifying activation overflow prevention at $x=15.0$, all matching types, specific node targeting, consecutive op cast elimination, creation op handling, and top-level cast_to_16_bit_precision forwarding.
   • Updated documentation in docs/src/utils/casting.md.
   • Added Towncrier changelog fragment changelog.d/7.added.

Verification

• All unit and end-to-end tests pass (106/106 in tests/casting/):

  pytest tests/casting/ -v

• All repository quality checks pass (26/26):

  make check

Closes #7