def METHOD_NAME(input_dim):
    hidden_dim = input_dim * 3
    output_dim_multiplier = input_dim - 4
    mask, _ = create_mask(
        input_dim, [hidden_dim], np.random.permutation(input_dim), output_dim_multiplier
    )
    init_random_params, masked_dense = serial(MaskedDense(mask[0]))
    rng_key = random.PRNGKey(0)
    batch_size = 4
    input_shape = (batch_size, input_dim)
    _, init_params = init_random_params(rng_key, input_shape)
    output = masked_dense(init_params, np.random.rand(*input_shape))
    assert output.shape == (batch_size, hidden_dim)