Vmap initialize model

[pharringtonp19]

Coming from jax, I was thinking I could initialize the parameters of a MLP by vmapping a function which initializes parameters for a single layer. Given that vmap does not exist in dex, my assumption is that I can do this with map. I don't understand the effect system in dex (yet!) so I can't tell from the signature of map whether this is right thing to do.

data Params (m: Nat) (n: Nat) = AsParams {weights : (Fin m => Fin n => Float) & bias: (Fin m => Float)}

def init_layer_params (key: Key) (m: Nat) (n: Nat) : Params m n = 
    [k1, k2] = split_key key 
    AsParams {
        weights = rand_mat m n rand k1, 
        bias = rand_vec m rand k2
    }

nodes_input = [12, 32, 32]
nodes_output = [32, 32, 1]
keys = split_key (new_key 1) :: Fin 3=>Key