JAX MLP activation list can be smaller than layer list (index misalignment)

[cpaniaguam]

In jax_mlp.py, self.layers is built with one entry per layer, but self.activation_funs filters out "linear" activations.

LANfactory/src/lanfactory/trainers/jax_mlp.py

Lines 97 to 101 in 9301939

	self.activation_funs = [
	self.activations_dict[activation]
	for activation in self.activations
	if (activation != "linear")
	]

That means self.activation_funs can be shorter than self.layers.
Later, forward indexes self.activation_funs[i] using the layer index, which can raise IndexError or apply the wrong activation when "linear" appears before the final layer.

def test_mlp_jax_with_linear_hidden_activation():
    model = JaxMLP(
        layer_sizes=[10, 10, 1],
        activations=["tanh", "linear", "linear"],
        train_output_type="logprob",
        train=True,
    )

    rng = jax.random.PRNGKey(0)
    test_input = jnp.ones((5, 6))
    params = model.init(rng, test_input)  # fails with IndexError in old pattern
    output = model.apply(params, test_input)

    assert output.shape == (5, 1)

Observed failure:
IndexError: tuple index out of range at jax_mlp.py

Expected:
Activation mapping should stay layer-aligned even when "linear" is used in hidden layers.

[linear-code]

HSSM-728