[jmschrei]

Six bugs on main, each fixed in its own commit with a test that fails without the fix. All six predate #120. They came up while reviewing the test suite file by file. The existing tests used inputs, or tolerances, that hid them.

Fixes

• ablate_annotations and marginalize_annotations on models with several outputs. Both built their list of outputs with range(len(y_befores)), which counts annotations rather than model outputs. On ConvDense, which has two outputs:

  • one annotation returned only the first output;
  • three annotations raised IndexError: list index out of range.

  test_ablate_annotations_multihead and test_marginalize_annotations_multi_output passed only because they used two annotations. Both now run one, two and three annotations.

• apply_pairwise required args in practice but defaulted it to None. Every call without args raised TypeError: 'NoneType' object is not iterable from inside the batching loop. args is now a required parameter, as it is in apply_product. Calls that pass args (the tutorial, the skill and every test do) are unchanged.

• characters and one_hot_to_fasta on numpy arrays. Both list numpy.ndarray in their signatures, but characters called .max(dim=0) on the array and raised TypeError: _amax() got an unexpected keyword argument 'dim'. one_hot_to_fasta failed the same way through characters. Both now convert an array to a tensor first, copying it only when it is not contiguous. As a result, one_hot_to_fasta also writes the all-zero columns of an array as N, as it does for a tensor.

• plot_logo on alphabets with characters beyond A, C, G and T. The default colors covered only A, C, G and T. A nonzero row for any other character, such as N or an amino acid, raised KeyError. With color=None, such characters are now drawn gray. A color dict that lacks a character still raises. test_plot_logo_alphabet_custom left its fifth row at zero, so it never drew N; it now draws it and checks the colors.

• The softmax rule of deep_lift_shap and pisa under bf16 autocast. The rule computed in the dtype of the cached activations, which is bf16 under dtype=torch.bfloat16.

  • bf16 holds the differences the rule divides to about three digits, and its 1e-6 guards are below bf16's resolution.
  • On ConvSoftmax, whose attributions reach 0.062, the bf16 attributions were off from float32 by up to 0.021 on the CPU and 0.030 on CUDA.
  • The rule now computes in at least float32, and _b_hook casts the multipliers back as before.
  • Float32 and float64 activations take the same path as before: .to() returns the tensor itself when the dtype already matches.

  Largest difference from float32, before → after:

  Model	CPU	CUDA
  ConvSoftmax	2.1e-2 → 2.2e-4	3.0e-2 → 2.7e-4
  ConvAttentionPool	7.1e-4 → 1.6e-4	1.2e-3 → 1.6e-4
  MultiHeadAttention	4.3e-3 → 6.7e-4	7.5e-4 → 7.6e-4
  TransformerBlock (2 blocks)	1.6e-2 → 1.6e-3	3.9e-3 → 1.9e-3

  The existing autocast test compared only ConvPoolDense, to two decimals. test_deep_lift_shap_autocast_softmax_matches_full_precision checks these four models: without the fix it fails in 7 of 8 cases, all except MultiHeadAttention on CUDA.

Checks

• Each new or changed test was run before its fix and failed with the error listed above.
• With the fixes, the full suite passes on the CPU and on one CUDA device: 3,572 passed, 9 skipped.
• sphinx -W passes.

The six surfaces

1. Tests: included in each commit.
2. Documentation:
   • The docstrings of characters and one_hot_to_fasta list numpy.ndarray.
   • plot_logo's color=None entry says that other characters are drawn gray.
   • No new page is needed.
3. README: no snippet uses these calls.
4. Changelog: one bullet per fix under Unreleased. The ablate and product subsections are new.
5. Bundled skill: no change needed.
   • references/motif-effects.md already says that multi-output models return lists.
   • references/product.md and references/func-pattern.md always pass args.
6. Compatibility: no model, state_dict or kernel is touched.

🤖 Generated with Claude Code