Applying Flax Model with Manually Modified Parameters

[mKabouri]

Hello,

I am trying to apply my flax model using modified parameters new_params. I generated new_params by manually doing some operations on the original parameters (state.params) using jax.tree.map. However, when I apply my model with new_params, the output type differs from the one I get when using the original state.params.

With state.params, I receive a JAX Array as the output. However, when using new_params directly (without involving the state), the output is a BatchTracer instead.

Basically what I want is doing a forward pass without updating the parameters of my network.

My function for that looks like that:

def apply_model_with_new_params(
    model: Network,
    new_params: Dict,
    input: jnp.ndarray
):
    return model.apply({'params': new_params}, input)

Could anyone explain why this is happening and how I can ensure consistent output types because I need JAX Array as the output and I do not want to update the state ?

Thanks!