[solarsys]

Problem

MLP and LogisticRegression fail on any timeseries input field:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (64x1008 and 21x128)

Both models flatten every 3D input to (batch, steps * width) before embedding. That is right for nested codes (nested_sequence), which become one code sequence. A numeric 3D input such as a time series (batch, time, features) is then flattened too, so it no longer matches the embedding's nn.Linear(features, dim). Above, 48 hourly steps of 21 lab and vital values became 1,008 columns. RNN, Transformer, CNN and RETAIN handle the same input, because they never flatten.

This came up running the in-hospital mortality task on the eICU demo with a 48-step lab and vital series.

Fix

Which inputs are embedded per step. A 3D input is embedded per step when its embedding is an nn.Linear: timeseries, nested_sequence_floats, stagenet_tensor and tensor, whose Linear runs over the last dimension. multihot is also embedded by an nn.Linear, but its samples are 1D, so it is not 3D in practice.

• The embedding projects each time step, and mean pooling averages the steps.
• Flattening exactly these inputs is what crashed, because the Linear takes one step's features. With a single step, the result is the same as before.
• Every other 3D input is flattened as on master. That includes nested codes (nested_sequence) and nested_multihot, which has no embedding layer and passes through as raw multi-hot rows.
• Tuple inputs in LogisticRegression. It now takes the value from a tuple input such as stagenet_tensor's (time, value), as MLP does. Before, it raised ValueError on any tuple input.

Padding.

• The mask comes from the raw values. Collate pads shorter series with all-zero steps, and the Linear maps each of them to its bias, which is nonzero. So for per-step inputs, the pooling mask is taken from the raw values with get_step_mask (pyhealth/models/utils.py): a step is padding when all its features are zero.
• Real all-zero steps count as padding too. A real step whose features are all zero, such as the zeros forward fill puts before the first measurement, is also treated as padding. The model docstrings say so.
• MLP.forward_from_embedding needs that mask passed in. It cannot tell padding from data in embedded values. Callers passing a Linear-embedded series must append the mask, as forward and the interpretability methods do. Its docstring says so.

mean_pooling, in both models.

• Masked steps are fully excluded. They are now left out of the sum as well as the count. Before, the mask only set the count.
• A row with no valid step no longer gives NaN. It is averaged over all its steps instead of dividing 0 by 0. This holds for every input type, nested codes included: a row whose steps are all padding now gets a finite value instead of NaN. For codes that value is a zero vector, because their padding embeds to zero.
• Other token rows are unchanged. Rows of token inputs that have at least one valid step give the same output as before.

A mask given with 3D codes. When a processor supplies its own mask, as tuple_time_text's attention mask does, MLP.forward flattens the mask with the codes and now passes the flattened mask on to the pooling. Masked tokens are therefore left out of the mean. A 3D mask that a caller of forward passes as an extra element after the schema's elements is flattened the same way. (The interpretability methods append their mask in this form, but for nested codes they go through forward_from_embedding, which does not support nested codes, so only direct callers of forward reach this path.) Before, the original 3D mask was passed on. On master that only skewed the count; with the new pooling, it would have raised.

Tests

In tests/core/test_mlp.py and tests/core/test_logistic_regression.py:

• test_timeseries_and_nested_codes: a timeseries field plus a nested_sequence field run a forward and backward pass. Without the fix, both models fail with the error above.
• test_other_numeric_3d_inputs_are_embedded_per_step: nested_sequence_floats, stagenet_tensor and tensor fields each run a forward and backward pass, and each sample pools to the mean of its embedded steps. The stagenet_tensor case fails on LogisticRegression without the tuple fix.
• test_padded_steps_are_left_out_of_the_mean: series of 2, 5 and 3 steps (the last all zeros) pool the same in one padded batch as each on its own. Each pooled vector equals the mean of the per-step embeddings of the real steps.
• test_nested_multihot_passes_through_flattened: documents the current passthrough of nested_multihot as a limitation. It passes on master too.
• test_all_zero_series_with_an_external_mask (MLP): forward_from_embedding with a raw-value mask matches forward for an all-zero series, and Integrated Gradients attributions are finite.
• test_mean_pooling_without_a_valid_step (LogisticRegression): checks the pooling rule directly.
• test_a_mask_given_with_3d_codes_is_flattened_with_them (MLP): nested codes with a mask that masks out a real code, given once as the processor's mask and once appended. Both pool to the mean of the kept codes' embeddings, computed by hand, and differ from the unmasked output. The appended case fails without the fix.
• Input-form tests: test_forward_from_embedding_inputs (MLP) and test_list_inputs_and_unsupported_ranks (LogisticRegression).

Results:

• The 20 core test files that use MLP or LogisticRegression pass: 276 tests, 5 skipped. They include the interpretability methods, calibration, early stopping and reproducibility tests.
• Coverage (CI, 5c2c548): mlp.py 84.14%, logistic_regression.py 82.56%, models/utils.py 84.38%.
• tools/check_pr_rules.py passes.
• With the repo's ruff settings, the model files have no new warnings. The new tests use naive datetimes, like the existing timeseries test (DTZ001).
• On the eICU demo task, both models now train on the time-series view. The comparison with a GRU (0.261 vs 0.240 cross-validated PR-AUC) was measured on the first commit, before the padding changes, and has not been rerun.

🤖 Generated with Claude Code

[github-actions]

Test Coverage

Total: 70.29% (18098 / 25749 lines)

Modified file	Coverage	Lines
pyhealth/autoimprove/__init__.py	100.00%	6 / 6
pyhealth/autoimprove/evaluate.py	92.07%	151 / 164
pyhealth/autoimprove/improver.py	95.63%	372 / 389
pyhealth/autoimprove/llm.py	90.62%	29 / 32
pyhealth/autoimprove/proposers.py	93.91%	108 / 115
pyhealth/autoimprove/space.py	95.31%	122 / 128
pyhealth/calib/predictionset/cluster/neighborhood_label.py	95.73%	112 / 117
pyhealth/datasets/base_dataset.py	90.80%	474 / 522
pyhealth/datasets/omop.py	100.00%	18 / 18
pyhealth/metrics/calibration.py	65.00%	52 / 80
pyhealth/metrics/multiclass.py	25.00%	26 / 104
pyhealth/models/__init__.py	100.00%	52 / 52
pyhealth/models/gradient_boosted_trees.py	96.79%	271 / 280
pyhealth/models/logistic_regression.py	82.56%	71 / 86
pyhealth/models/micron.py	83.23%	129 / 155
pyhealth/models/mlp.py	84.14%	122 / 145
pyhealth/models/survival_mdn.py	93.83%	213 / 227
pyhealth/models/utils.py	84.38%	27 / 32
pyhealth/processors/deep_nested_sequence_processor.py	91.37%	180 / 197
pyhealth/processors/nested_sequence_processor.py	93.53%	159 / 170
pyhealth/processors/sequence_processor.py	95.77%	68 / 71
pyhealth/processors/stagenet_processor.py	84.08%	169 / 201
pyhealth/tasks/cardiology_detect.py	93.68%	89 / 95

Workflow run · 5c2c548

[Logiquo]

/review

[github-actions]

Previous review

AI review: findings

The PR stops MLP and LogisticRegression from flattening numeric 3D inputs, so a (batch, time, features) timeseries is embedded per step by nn.Linear instead of crashing. Token inputs (processor.is_token()) are still flattened, so nested_sequence behaves as before. Changed files are pyhealth/models/mlp.py, pyhealth/models/logistic_regression.py and one new test in each model's test file. Blocking: when series in a batch have different lengths, collate pads the shorter ones with zeros (pyhealth/datasets/utils.py:285-320), and the models then derive the pooling mask from the embedded values. Each Linear has a bias, so padded steps embed to a nonzero vector, count as valid, and dilute the mean pooling; this is silent and affects the eICU time-series path the PR targets (mlp.py:357, logistic_regression.py:181). Non-blocking: the new tests use equal-length series and only assert output shape, so they miss this and do not check per-step embedding; they do cover the original crash. I found no existing component that duplicates this change. I could not run the test suite, so these conclusions come from reading the code. Overall risk is medium: the crash fix is sound, but the pooling bug should be fixed or tested before merge.

Findings outside the diff

• [high] Zero-padded time steps are counted in LogisticRegression mean pooling pyhealth/models/logistic_regression.py:181
  Same problem as in MLP. LogisticRegression.forward now passes an unflattened numeric series to the embedding. The pooling mask at line 181, (x.sum(dim=-1) != 0), is computed on the embedded tensor, so collate zero-padding (datasets/utils.py:285-320) embeds to the nonzero bias and is counted as valid. Fix: derive the mask from the raw input before embedding, and cover variable-length series in the test.

_{Reviewed commit 125689a}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: need change

Thank you for this contribution. Embedding time series per step in MLP and LogisticRegression, instead of flattening them, fixes a real crash, and the regression tests for it are welcome.

Coverage: the whole repository is at 67.94%, which meets the 60% bar. The changed files are below the 80% guideline: logistic_regression.py is at 75.68% and mlp.py is at 75.19%. I'm not asking for more coverage on its own, but the test suggested below should raise it.

To get this merged:

1. Handle padding for variable-length series. When series in a batch have different lengths, collate pads the shorter ones with zeros. The models then derive the pooling mask from the embedded values (mlp.py around line 357, logistic_regression.py around line 181). Both Linear layers have a bias, so a padded step embeds to a nonzero vector. It is counted as valid and dilutes the mean pooling, silently. This hits the eICU time-series path the PR targets. Please compute the mask from the raw input before embedding, for example a step is padding if all its features are zero. Or use a mask from the processor if one exists.
2. Add a test with series of different lengths in one batch. It should check that padding does not change the pooled result. For example, compare a short series in a padded batch with the same series on its own. The current tests use equal-length series and only assert the output shape, so they would not catch this.

Optional suggestions:

• Make the new tests assert something about per-step embedding beyond the output shape.

_{Reviewed commit 125689a}

[solarsys]

Thanks for the review. The padding issue is fixed in 6f84180. test_mlp.py (7 tests) and test_logistic_regression.py (8 tests) pass with the CI command, and so do the 16 other test files that use these models (252 tests in all, including the interpretability tests). The PR rules check passes.

Required

1. Padded steps are left out of the mean. There were two parts to this:
   • The mask came from the embedded values. The new get_step_mask (pyhealth/models/utils.py) builds it from the raw series instead: a step is padding when all its features are zero.
     • MLP.forward passes this mask to forward_from_embedding as the extra mask element it already accepts, so that method's signature is unchanged. LogisticRegression uses it for its pooling.
     • Token inputs keep the mask from the embedded values, which was already right because their padding embeds to zero. A mask from the processor schema is used as before.
     • A series with no nonzero step keeps all its steps, so pooling it never divides by zero.
   • mean_pooling ignored the mask in the numerator. It summed every step and used the mask only in the count (x.sum(dim=1) / mask.sum(...)). Even a correct mask would have left the padded steps' bias in the mean. Both models now multiply by the mask before summing. Token outputs are unchanged, because masked token steps are zero vectors. For the one processor that supplies a schema mask (tuple_time_text), padded notes now drop out of the numerator as well.
2. A test with series of different lengths. test_padded_steps_are_left_out_of_the_mean in both test files puts series of 2, 5 and 3 steps (the last all zeros) in one batch. It checks that each sample's pooled embedding equals the same sample run on its own. Without the fix, the test fails for both models; the largest difference is 0.18 for MLP and 0.74 for LogisticRegression.

Optional

• Per-step embedding. The same test also checks the pooled vector against the mean of the embedding's nn.Linear applied to each real step: through model.mlp for MLP, and directly for LogisticRegression. The original test_timeseries_and_nested_codes stays as the regression test for the crash.

One limitation remains. A real step whose features are all zero cannot be told apart from padding, because the timeseries processor has no mask. Such a step is now left out of the mean, both in a padded batch and on its own.

[solarsys]

/review

[github-actions]

Previous review

AI review: findings

The PR makes MLP and LogisticRegression embed numeric 3D inputs per time step instead of flattening them, then mean-pool over the real steps using a mask taken from the raw values. Token inputs keep the flatten path, and the reported timeseries crash is fixed. The change touches only pyhealth/models/mlp.py and logistic_regression.py, adds get_step_mask to models/utils.py, and changes no base class or shared default. The public forward_from_embedding contract is unchanged, and the interpretability methods (IG, GIM, SHAP, LIME, DeepLift) already append raw masks, so they stay consistent. Redundancy: no existing component does this for these two models, and the inline mask expressions in rnn/tcn/stagenet are not reusable helpers, so there is no redundancy finding. Blocking (medium), mlp.py:364: the is_token() gate also catches NestedMultiHotProcessor (is_token False, no embedding layer, per-visit categorical vectors), which was flattened before, so the PR's claim that nested codes behave as before is wrong for it and nothing tests it. Blocking (medium), mlp.py:180: mean_pooling divides by a zero mask sum for an all-zero series when an interpretability method passes its raw-value mask, so base logits and attributions become NaN while plain forward stays finite, and the all-zero branch of get_step_mask is untested. Non-blocking (low), utils.py:70: any all-zero step counts as padding, so real zero steps such as leading forward-fill zeros are dropped from pooling, which the docstring states. The new tests cover the reported crash, padded-batch equivalence against single-sample runs, and a backward pass, matching the repo's usual level, though they check only shapes and equality. The 0.261 vs 0.240 benchmark and the 20-file test claim cannot be checked from the repository. Overall risk is low for the reported path, with the two blocking items affecting edge inputs.

_{Reviewed commit 6f84180}

[github-actions]

Previous review

AI Review

suggestion: need change
decision: need change

Thank you for this contribution. Embedding time series per step in MLP and LogisticRegression, instead of flattening them, fixes a real crash, and the padding-aware pooling and the tests are well thought out. It is close, but there are two blocking problems.

To get this merged:

1. The not processor.is_token() check in mlp.py (around line 364) and in logistic_regression.py also matches processors that are not token-based but whose 3D inputs were flattened before. One example is NestedMultiHotProcessor, which gives per-visit categorical vectors. These inputs now take the new per-step path, so the claim that nested codes behave as before does not hold for them. Please narrow the condition so that only true numeric series such as timeseries take the new path. Alternatively, keep the old flatten behaviour for every other 3D input. Please also add a test for a multi-hot nested input.
2. In MLP, mean_pooling (around line 180) divides by mask.sum(dim=1). When an interpretability method supplies its own raw-value mask for an all-zero series, that sum is 0. The result is NaN base logits and attributions, even though a plain forward pass stays finite. Please guard against a zero mask sum, for example by clamping the denominator to a minimum of 1. Please also add a test for an all-zero series, covering both the get_step_mask fallback and the case where an external mask is passed.

Optional suggestions:

• Treating every all-zero step as padding drops real zero steps, such as leading forward-filled zeros, from the pooling. Your docstring says this, but it would help to note it in the model docs as well.
• Coverage is 69.11% for the whole repository, which meets the 60% bar. Several changed files are below the 80% target: mlp.py at 76.47% and logistic_regression.py at 77.78%. Tests for the two items above should help raise them. utils.py is at 85.29%.

_{Reviewed commit 6f84180}

[solarsys]

Thanks for the second review. Both blocking items are fixed in f21e45e.

• Tests: test_mlp.py and test_logistic_regression.py have 11 tests each, and both pass with the CI command. So do the 16 other test files that use these models: 259 tests in all, 5 skipped, including the interpretability tests.
• Coverage: measured locally over those files, mlp.py is at 86.23% and logistic_regression.py at 81.93%.
• PR rules check: passes.

Blocking

1. Only real numeric series take the per-step path. I checked every registered processor's is_token(), output rank and the embedding EmbeddingModel creates for it.
   • The rule: a 3D input is embedded per step only when its embedding is an nn.Linear over the last dimension. That covers timeseries, nested_sequence_floats, stagenet_tensor and tensor. Flattening exactly these inputs crashed, because the Linear takes one step's features. With a single step, the per-step result equals the old flattened one.
   • Everything else is flattened as on master. This includes nested_multihot, which is not a token input but has no embedding layer. It also covers any processor without a Linear embedding, such as temporal_timeseries.
   • Test: test_nested_multihot_is_flattened, in both test files, checks that the pooled vector is the flattened multi-hot batch: through model.mlp for MLP, and directly for LogisticRegression. It passes against master's model code and fails on 6f84180 with a shape error.
2. No NaN for an all-zero series.
   • The fix: mean_pooling, in both models, averages a row with no valid step over all its steps, the same as having no mask, and clamps the count at 1. A raw-value mask for an all-zero series, built the way the interpretability methods build it, now gives finite output. It also matches forward, because forward applies the same rule to that row. get_step_mask now only marks the steps that hold data, so the all-zero rule lives in one place.
   • test_padded_steps_are_left_out_of_the_mean (both files) runs an all-zero series through forward and asserts finite output that matches the series run on its own.
   • test_all_zero_series_with_an_external_mask (MLP) passes that mask to forward_from_embedding and checks the logits equal forward's. It also checks that Integrated Gradients attributions are finite.
   • test_mean_pooling_without_a_valid_step (LogisticRegression) checks the pooling rule directly.
   • On 6f84180: the forward-from-embedding test and the pooling test fail.

Optional

• Docs. The MLP and LogisticRegression class docstrings now say that an all-zero step of a numeric series counts as padding, for example the zeros forward fill puts before the first measurement. They also say that a series with no nonzero step is averaged over all its steps.

• Coverage. I added small tests for lines that were cheap to reach:

  • MLP: the constructor's argument checks and activations, sum_pooling, and forward_from_embedding with a bare tensor, with a mask that has a feature dimension, and with an unsupported rank.
  • LogisticRegression: inputs given as lists, and an unsupported rank.

  The lines still uncovered in mlp.py are mostly the __main__ demo and the path for processors that supply their own mask, which needs a transformers text model.

[solarsys]

/review

[github-actions]

Previous review

AI review: findings

The PR stops MLP and LogisticRegression from flattening 3D numeric inputs that an nn.Linear embeds (for example timeseries). Each time step is now embedded and mean-pooled over steps with nonzero features, using a mask from the new get_step_mask in pyhealth/models/utils.py. Token inputs such as nested codes are still flattened, and mean_pooling in both models now averages all steps for a row with no valid step instead of returning NaN. The interpretability methods already append raw-value masks, so their results match forward(). Blocking problems: none found. Non-blocking: the description says the split uses processor.is_token(), but the code checks for nn.Linear; the two agree for current processors. Non-blocking: real all-zero steps are excluded from pooling (documented), and forward_from_embedding without a mask still counts padded steps, unlike forward(). Non-blocking: mean_pooling and the per-step logic are copied between mlp.py and logistic_regression.py, the two test helpers are copied between the test files, and one test pins nested_multihot's passthrough. The new tests cover shape and backward pass, padded-batch pooling equal to single-series pooling, all-zero rows, and an IG run with a mask, in the style of the existing MLP and LR tests. I could not run the test suite or the eICU experiment with the tools available here, so the claimed pass counts and PR-AUC figures are unverified. Overall risk is low: behavior changes only for numeric 3D inputs that previously crashed and for all-invalid rows that previously gave NaN.

Findings outside the diff

• [low] forward_from_embedding without a mask counts padded steps, unlike forward pyhealth/models/mlp.py:279
  When no mask is passed, the mask is read from the embedded values. For a Linear-embedded series, padded steps embed to the bias, which is nonzero, so they are counted. forward() excludes them, so the same batch gives different logits through forward() and through forward_from_embedding(bare tensor). The interpretability methods append masks and are not affected, but direct callers are. Fix: document this in forward_from_embedding, or build get_step_mask(value) for 3D Linear-embedded inputs when no mask is given.

_{Reviewed commit f21e45e}

[github-actions]

Previous review

AI Review

suggestion: human review
decision: not available

Haiku found that this change needs a human review: The change alters the forward_from_embedding contract that every interpretability method (IG, SHAP, LIME, DeepLift, GIM, IG-GIM) calls. It adds a trailing-mask convention and changes pooling semantics for numeric 3D inputs and all-invalid rows, and it adds a function to the shared pyhealth/models/utils.py. It does not edit a base class, but the reach into interpretability code means a human should confirm it.

_{Reviewed commit f21e45e}

[solarsys]

Thanks for the third review. The low findings are addressed in 3887d2c and 5f0ee46, and I've left a few as they are, with reasons below.

• Tests: test_mlp.py and test_logistic_regression.py have 11 tests each, and both pass with the CI command. So do the other files that use these models: 18 files and 259 tests in all, 5 skipped.
• Coverage: measured locally over those files, mlp.py is at 83.57% and logistic_regression.py at 81.93%.
• PR rules check: passes.
• Description: I've rewritten the PR description to match the code.

Changed

• forward_from_embedding without a mask counts padded steps. This is now documented in its docstring rather than fixed automatically. Embedded values can't tell padding from data, because a padded step embeds to the Linear's bias.
  • Why not infer the mask: the only automatic rule would be "a step equal to the bias is padding", and it isn't exact. The interpretability methods pass scaled or interpolated embeddings, in which a padded step is α·bias. A nonzero raw step in the Linear's null space would also embed to the bias.
  • What the docstring says: callers passing a Linear-embedded series must append the step mask, as forward and the interpretability methods do. The mask is get_step_mask of the raw values; the interpretability methods' abs().sum(-1) != 0 is equivalent.
  • LogisticRegression has no forward_from_embedding, so nothing changes there.
• Split rule in the description. The description now states the rule the code uses. A 3D input is embedded per step only when its embedding is an nn.Linear (timeseries, nested_sequence_floats, stagenet_tensor, tensor). Every other 3D input is flattened as on master. The processor.is_token() wording is gone.
• All-padding rows no longer give NaN, nested codes included. The description no longer says nested codes behave exactly as before. It now says that a row whose steps are all padding, for every input type, gets a finite value (a zero vector for codes) instead of NaN. Rows with at least one valid step are unchanged for token inputs.
• The nested_multihot test. It is renamed test_nested_multihot_passes_through_flattened in both files, with a comment that it pins the current passthrough, not a design choice. EmbeddingModel has no layer for nested_multihot, so the raw rows pass through and are flattened.
• The constructor test. test_arguments_it_sets_itself_and_activations is dropped as asked, since it checked code this PR does not touch. That took mlp.py to 79.71%. The test below brings it back to 83.57%.
• A processor's own mask with 3D codes. The new test_a_processor_mask_is_flattened_with_its_codes covers the flatten path when a processor supplies its own mask, as tuple_time_text's attention mask does. The processor's schema is patched to supply a mask over nested codes, and one real code per row is masked out. Each pooled vector must equal the mean of the kept codes' embeddings, computed by hand, and must differ from the unmasked output.
  • It found a bug, now fixed in 5f0ee46. forward flattened such a mask together with the codes, but passed the original 3D mask on to the pooling. On master that only skewed the count. With mean_pooling now multiplying by the mask, the shapes no longer matched and forward raised. forward now writes the flattened mask back, so masked tokens are left out of the mean.
  • Correction: this also corrects my earlier reply, which said padded tuple_time_text notes already dropped out of the sum. That is only true from 5f0ee46.

Left as they are

• Shared mean_pooling and step logic. The duplication between MLP and LogisticRegression predates this PR. Moving it into pyhealth/models/utils.py would widen the change to shared utilities, which is already what sent the PR to human review, so I've kept the edit local to the two models.
• Test helpers. _padded_series and _nested_multihot stay in each test file. The two files are otherwise independent, and a shared helper module would couple them for about 40 lines.
• An observed-flag mask. No current numeric processor supplies one, and the all-zero-step limitation is documented in get_step_mask and in both model docstrings. In MLP, a mask that a processor supplies in its schema is still used, and is now flattened with 3D codes as above.

[solarsys]

/review

[github-actions]

Previous review

AI review: findings

The PR makes MLP and LogisticRegression embed 3D numeric inputs (timeseries, nested_sequence_floats, stagenet_tensor, tensor) once per time step through their nn.Linear instead of flattening them, which fixes the timeseries shape crash. Pooling uses a mask built from raw values by a new get_step_mask helper in models/utils.py. mean_pooling in both models now excludes masked steps and returns finite values for all-masked rows. This changes the output of the public static mean_pooling methods and the mask contract of MLP.forward_from_embedding, which IG, SHAP, LIME and DeepLift call. I checked that those methods append raw-value masks, so timeseries stays consistent across them. Blocking (medium): tests/core/test_mlp.py:191 and the matching LogisticRegression test cover only timeseries, so the per-step path for tensor, nested_sequence_floats and the (time, value)-tuple stagenet_tensor form, which the description claims, has no test. Non-blocking (low): a 3D mask appended by interpretability methods for nested codes is not flattened in MLP.forward (mlp.py:190), so the new pooling raises a broadcast error where master miscounted; 3D MultiHot inputs are now embedded per step, contrary to the description; the nested_multihot test pins a passthrough; and the test helpers are duplicated across two files. The padding and pooling tests are thorough, and the eICU numbers are stale as the PR says. I could not run the test suite with read-only tools, so the 259-test pass claim is unverified. Overall risk is low, with one blocking test gap.

_{Reviewed commit 5f0ee46}

[github-actions]

Previous review

AI Review

suggestion: human review
decision: not available

Haiku found that this change needs a human review: The change edits MLP.forward_from_embedding and the public static mean_pooling methods of MLP and LogisticRegression. forward_from_embedding is the embedding-level contract that IG, SHAP, LIME, DeepLift and GIM call, and its mask semantics and pooling outputs change. It also adds a helper to the shared pyhealth/models/utils.py module.

_{Reviewed commit 5f0ee46}

[solarsys]

Thanks for the review. Fixed in 5c2c548.

• Per-step path tested only for timeseries (blocking). test_other_numeric_3d_inputs_are_embedded_per_step, in both test files, runs nested_sequence_floats, stagenet_tensor and tensor fields forward and backward. It also checks that each sample pools to the mean of its embedded steps. The test found that LogisticRegression raised on stagenet_tensor's (time, value) tuple, as it does on master. It now takes the value from the tuple, as MLP does.
• Appended 3D mask with nested codes. MLP.forward now flattens it with the codes, as it does a processor's mask. test_a_mask_given_with_3d_codes_is_flattened_with_them covers both kinds of mask. The appended case failed before the fix with the broadcast error you describe.
• Per-step condition and multi-hot. The description now says multihot is also embedded by an nn.Linear but is not 3D in practice. I kept the check on the embedding type rather than a list of processors, so a new numeric processor with a Linear embedding gets the per-step path without another edit.
• nested_multihot test. Its docstring now says it documents a limitation.
• Left as before, for the reasons in my last reply: the duplicated test helpers, and all-zero steps counting as padding, which the docstrings state. The processor-mask case still patches the schema, but the same test now also runs the unpatched appended-mask path.

The 20 core test files that use MLP or LogisticRegression pass (276 tests, 5 skipped), as does tools/check_pr_rules.py. CI coverage on 5c2c548: mlp.py 84.14%, logistic_regression.py 82.56%, models/utils.py 84.38%.

[solarsys]

/review

[github-actions]

AI review: findings

The PR embeds 3D numeric inputs (timeseries, nested_sequence_floats, stagenet_tensor, tensor) step by step in MLP and LogisticRegression instead of flattening them, which fixes the crash on timeseries input. Padded steps are masked from raw values with a new get_step_mask in pyhealth/models/utils.py, mean_pooling excludes masked steps and no longer returns NaN, a processor or appended 3D mask is flattened with 3D codes, and LogisticRegression reads the value from tuple inputs. Interactions: the change stays inside the two models plus an additive utility, but it changes outputs of MLP.forward, MLP.forward_from_embedding and the static mean_pooling in both models, which the interpretability methods (IG, DeepLift, SHAP, LIME, GIM) call, so IG attributions on padded timeseries MLP inputs and outputs for masked token rows change. Redundancy: nothing existing does this, though the one-line step-mask rule overlaps EHRMamba's _create_mask, and mean_pooling and the per-step block are duplicated across MLP and LogisticRegression (non-blocking, low). Blocking problems: none found. Non-blocking: the nested_multihot test enshrines a passthrough that only works when the visit count matches embedding_dim, and the description overstates that interpretability methods reach the appended 3D-mask path. Tests cover the new paths well (per-step embedding, padding, masks, tuple input, IG finiteness), but I did not execute them because no shell was available, so the reported pass counts and coverage are unverified. Real all-zero steps are dropped as padding, which the docstrings state. Risk: low.

_{Reviewed commit 5c2c548}

[github-actions]

[low] mean_pooling and the per-step block are copied from MLP

The new mean_pooling body (lines 127-129) is the same as MLP.mean_pooling (pyhealth/models/mlp.py:188-190), and the per-step detection and mask code in LogisticRegression.forward (lines 168-178) mirror MLP.forward (mlp.py:375-387). The duplication predates this PR, but the PR edits both copies in parallel and they can drift. Consider calling MLP.mean_pooling from LogisticRegression or moving the shared pooling into models/utils.py. Not blocking.

[github-actions]

[low] get_step_mask repeats an inline rule that already exists in EHRMamba

get_step_mask computes (x != 0).any(dim=-1), which is the same rule as EHRMamba._create_mask for TimeseriesProcessor (pyhealth/models/ehrmamba.py:186, torch.any(torch.abs(value) > 0, dim=-1)). The interpretability methods also inline the same expression (integrated_gradients.py:316 and others). It is not a full duplicate, since no existing helper does this job for MLP and LogisticRegression, so this is not blocking. Optionally reuse it in EHRMamba or the interpretability methods.

[github-actions]

[low] nested_multihot test uses a fixed visit count that hides a crash

The passthrough test (and its MLP twin at tests/core/test_mlp.py:310) gives every sample two visits, so the flattened width always equals embedding_dim. With variable visit counts, collate pads the visit dimension to the batch maximum, the flattened width changes per batch, and the fixed nn.Linear (MLP.mlp and LogisticRegression.fc) fails. The PR states the limitation, so this is not blocking. Consider adding a batch with differing visit counts so the limitation is shown, or raising a clear error for nested_multihot.

[github-actions]

[low] Appended 3D-mask flattening is not reached by the interpretability methods as described

The description says the appended 3D mask is flattened 'the way the interpretability methods append theirs'. For nested codes, IG embeds the codes with nn.Embedding (integrated_gradients.py:366-368), which gives 4D values. forward_from_embedding rejects 4D values (mlp.py:298-301), so IG never reaches this branch, which only direct callers of forward use. The branch is correct for direct callers, but the description should say so, or the test should cover forward_from_embedding with a flattened input. Not blocking.

[github-actions]

AI Review

suggestion: human review
decision: not available

Haiku found that this change needs a human review: The change alters existing public behavior of MLP and LogisticRegression, not just new code paths. MLP.forward_from_embedding and the static mean_pooling in both models now return different values for masked rows, and interpretability methods (IG, DeepLift, SHAP, LIME, GIM) call forward_from_embedding, so their attributions on padded timeseries MLP inputs change. It also adds a helper to the shared pyhealth/models/utils.py. These are output contracts other code relies on, so a human should review them.

_{Reviewed commit 5c2c548}

[solarsys]

Thanks. The review found no blocking problems. Here is how I've handled the four lows:

• mean_pooling and the per-step block are copied from MLP. The duplication was already there before this PR. I kept it so that the PR only changes the two forward paths. Moving the pooling into pyhealth/models/utils.py would be a good follow-up.
• get_step_mask and EHRMamba. The rule is the same, but switching EHRMamba and the interpretability methods to the helper would change code outside this PR, so I left them as they are.
• The nested_multihot test. Agreed: when visit counts differ across batches, the flattened width changes and the Linear layer fails. Master behaves the same way. The test only records the current passthrough, and the description calls it a limitation. Handling nested_multihot properly, or raising a clear error, belongs in its own PR.
• The appended 3D mask. I edited the description after the review. It now says that only direct callers of forward reach this path.

I also merged this branch into current master locally (master is 22 commits ahead), and the MLP and LogisticRegression tests pass there.