Keep JAX MLP activations aligned with layers for hidden "linear"

[fmuia]

Closes #105.

Problem

JaxMLP.setup() dropped "linear" entries from activation_funs, but __call__ indexes that list by layer position. So any "linear" before the output layer moved every later activation one layer earlier:

• With one activation per layer (the normal case), this always ends in IndexError: tuple index out of range, e.g. ["tanh", "linear", "linear"].
• With an activations list longer than layer_sizes, it runs but applies the wrong activation. For example, ["linear", "tanh", "linear", "tanh"] on 3 layers applies tanh to layer 0.

jaxtrain can hit this: _get_train_network_config appends the output activation to the YAML's ACTIVATIONS, so ACTIVATIONS: [["tanh", "linear"]] crashes.

Fix

Following the recipe in the issue thread: map "linear" to an identity in activations_dict, keep one entry per activation in activation_funs, and collapse the forward loop to x = self.activation_funs[i](lyr(x)). This also removes the now-obsolete TODO: Warn if linear activation is used before final layer.

Two deliberate choices:

• lambda x: x, not jax.nn.identity. jax.nn.identity is jitted, so it would add a nested jit call to the traced graph. The lambda adds nothing, so the exported ONNX for existing configs is unchanged.
• Indexing, not zip. Plain zip would silently drop layers if activations were too short. zip(strict=True) would reject over-long lists that work today (a YAML that already lists the output activation gets a second one appended by the CLI). Indexing keeps current behaviour for both.

Tests

test_mlp_jax_activations_stay_layer_aligned runs 4 configs: the default plus three with a hidden "linear". It compares each output with the per-layer maths computed from the initialized params, because a shape check alone can't tell whether each layer got its own activation. On main the three hidden-"linear" cases fail with IndexError; with the fix all four pass.

Also checked locally:

• All 60 tests in the files that use JaxMLP or jaxtrain pass: test_jax_mlp.py, onnx/test_jax_export.py, test_end_to_end_jax.py, test_jax_lr_dict_else.py, test_run_identity.py, test_mlflow_integration.py, cli/test_cli_entry_points.py.
• ruff check src/lanfactory and ruff format --check . are clean. mypy reports the same pre-existing errors as on main, none in changed code.
• Default LAN and CPN ([100, 100, 1], tanh, tanh, linear): the traced jaxpr and the exported ONNX are byte-identical before and after. Parameter names are unchanged, so existing *_train_state.jax files still load.
• A CPN with a hidden "linear" layer exports, passes assert_single_trial_contract, and onnxruntime matches JAX.

Not in this PR

• TorchMLP rejects the same config with KeyError: 'linear' because its activation dict has no "linear". Adding "linear": torch.nn.Identity() would make the two backends match; left for a follow-up.
• A len(activations) == len(layer_sizes) check would catch the over-long case, but it would break YAMLs that already list the output activation, so it needs a decision first.

🤖 Generated with Claude Code

https://claude.ai/code/session_01LRVhzVTfUsMkPu83XJbh8E

---

Generated by Claude Code

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
src/lanfactory/trainers/jax_mlp.py	96.20% <100.00%> (-0.09%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.