Skip to content

Fix dtype upcast in removal-based ablation for discrete features#1170

Open
fbonc wants to merge 1 commit into
sunlabuiuc:masterfrom
fbonc:fix-ablation-dtype
Open

Fix dtype upcast in removal-based ablation for discrete features#1170
fbonc wants to merge 1 commit into
sunlabuiuc:masterfrom
fbonc:fix-ablation-dtype

Conversation

@fbonc

@fbonc fbonc commented Jun 28, 2026

Copy link
Copy Markdown
Contributor

Contributor: Felipe Amaral Bonchristiano (felipea5@illinois.edu)

Contribution Type: Fix dtype bug in base.py (metric-removal ablation)

Description

The plain-tensor branch of RemovalBasedMetric._apply_ablation (in pyhealth/metrics/interpretability/base.py) ablates discrete code-index inputs by multiplying them with a float mask:

ablated_inputs[key] = x * (1 - mask)   # x is Long, mask is Float -> result is Float

This upcasts integer code indices to float. The downstream nn.Embedding then raises:

RuntimeError: Expected tensor for argument #1 'indices' to have scalar type Long or Int; but got torch.FloatTensor instead

so comprehensiveness/sufficiency evaluation crashes for any plain-tensor (non-tuple) discrete input. e.g. a vanilla Transformer. The tuple-input branch (StageNet-style (time, values)) already guarded against this with a dtype check + integer-preserving multiply; the plain-tensor branch did not.

Hit while building a MIMIC-IV-demo benchmark for the attention-rollout method (#1158). Running evaluate_attribution with comprehensiveness/sufficiency on a plain Transformer triggered the crash. This is a pre-existing bug independent of that PR.

Fix

Detect discrete (integer) inputs and zero-ablate them to the padding index with (1 - mask).long(), which preserves dtype. To avoid duplicating that logic across the tuple and plain-tensor branches, the dtype check, integer ablation, the continuous zero/mean/noise strategies, and the empty-row safety repair are extracted into a single shared _apply_ablation_to_values helper that both branches call. Behavior for continuous inputs is unchanged.

Discrete features are always zero-ablated to the padding index regardless of ablation_strategy (mean/noise are meaningless for categorical code indices); this is now documented in the helper.

Files to review:

Base.py: bug fix

The plain-tensor branch of RemovalBasedMetric._apply_ablation multiplied
integer code-index inputs by a float mask (x * (1 - mask)), upcasting
them to float. Downstream nn.Embedding then raised
"Expected ... Long/Int but got FloatTensor", crashing comprehensiveness/
sufficiency evaluation on plain Transformer models. The tuple-input
branch (e.g. StageNet) already handled this; the plain-tensor branch did
not.

Detect discrete (integer) inputs and zero-ablate them to the padding
index with an integer-preserving (1 - mask).long() multiply. The dtype
check, integer ablation, and the empty-row safety repair are extracted
into a shared _apply_ablation_to_values helper so both input formats
follow one path.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant