Get module parameters after dropout has been applied

[bstaber]

Hello,

I'm using a simple feedforward network with dropout layers in order to perform Monte Carlo dropout. At inference time, I call the network as follows:

init_key = jax.random.PRNGKey(0)
rng_keys = jax.random.split(init_key, 2000)

@jax.jit
def predict_fn(dropout_rng):
    return jax.vmap(lambda x: network().apply({"params": state.params}, x, rngs={"dropout": dropout_rng}, deterministic=False), 0)(X_test)

predictions = []
for rng_key in tqdm(keys):
    prediction = predict_fn(rng_key)
    predictions += [prediction.squeeze()]
predictions = jnp.stack(predictions, axis=0)

In addition to predictions, I would like to get the values of parameters after the random masks have been applied in the dropout layers. Is this possible ?

Thank you !

[cgarciae]

Hey @bstaber, check out the Extracting intermediate values guide.