Flax Key -> Haiku Key Mapping
--------------------------------------------------------------------------------
params/a                                                     -> alpha_net_hiku/a
params/b                                                     -> alpha_net_hiku/b
params/ftes_0/vec_proj/kernel                                -> alpha_net_hiku/~/fte/~/linear/w
params/ftes_0/xvec_proj/layers_0/bias                        -> alpha_net_hiku/~/fte/~/linear_1/b
params/ftes_0/xvec_proj/layers_0/kernel                      -> alpha_net_hiku/~/fte/~/linear_1/w
params/ftes_0/xvec_proj/layers_2/bias                        -> alpha_net_hiku/~/fte/~/linear_2/b
params/ftes_0/xvec_proj/layers_2/kernel                      -> alpha_net_hiku/~/fte/~/linear_2/w
params/ftes_1/vec_proj/kernel                                -> alpha_net_hiku/~/fte_1/~/linear/w
params/ftes_1/xvec_proj/layers_0/bias                        -> alpha_net_hiku/~/fte_1/~/linear_1/b
params/ftes_1/xvec_proj/layers_0/kernel                      -> alpha_net_hiku/~/fte_1/~/linear_1/w
params/ftes_1/xvec_proj/layers_2/bias                        -> alpha_net_hiku/~/fte_1/~/linear_2/b
params/ftes_1/xvec_proj/layers_2/kernel                      -> alpha_net_hiku/~/fte_1/~/linear_2/w
params/ftes_2/vec_proj/kernel                                -> alpha_net_hiku/~/fte_2/~/linear/w
params/ftes_2/xvec_proj/layers_0/bias                        -> alpha_net_hiku/~/fte_2/~/linear_1/b
params/ftes_2/xvec_proj/layers_0/kernel                      -> alpha_net_hiku/~/fte_2/~/linear_1/w
params/ftes_2/xvec_proj/layers_2/bias                        -> alpha_net_hiku/~/fte_2/~/linear_2/b
params/ftes_2/xvec_proj/layers_2/kernel                      -> alpha_net_hiku/~/fte_2/~/linear_2/w
params/ftes_3/vec_proj/kernel                                -> alpha_net_hiku/~/fte_3/~/linear/w
params/ftes_3/xvec_proj/layers_0/bias                        -> alpha_net_hiku/~/fte_3/~/linear_1/b
params/ftes_3/xvec_proj/layers_0/kernel                      -> alpha_net_hiku/~/fte_3/~/linear_1/w
params/ftes_3/xvec_proj/layers_2/bias                        -> alpha_net_hiku/~/fte_3/~/linear_2/b
params/ftes_3/xvec_proj/layers_2/kernel                      -> alpha_net_hiku/~/fte_3/~/linear_2/w
params/kernel1                                               -> alpha_net_hiku/kernel1
params/kernels_imag                                          -> alpha_net_hiku/kernels_imag
params/kernels_real                                          -> alpha_net_hiku/kernels_real
params/last_layer/bias                                       -> alpha_net_hiku/~/linear_4/b
params/last_layer/kernel                                     -> alpha_net_hiku/~/linear_4/w
params/last_layer_quantum/bias                               -> alpha_net_hiku/~/linear_5/b
params/last_layer_quantum/kernel                             -> alpha_net_hiku/~/linear_5/w
params/lin/layers_0/bias                                     -> alpha_net_hiku/~/linear_2/b
params/lin/layers_0/kernel                                   -> alpha_net_hiku/~/linear_2/w
params/lin/layers_2/bias                                     -> alpha_net_hiku/~/linear_3/b
params/lin/layers_2/kernel                                   -> alpha_net_hiku/~/linear_3/w
params/message_layers_0/dia/bias                             -> alpha_net_hiku/~/equi_message_passing/~message/linear_1/b
params/message_layers_0/dia/kernel                           -> alpha_net_hiku/~/equi_message_passing/~message/linear_1/w
params/message_layers_0/diachi1                              -> alpha_net_hiku/~/equi_message_passing/diachi1
params/message_layers_0/diagonal/layers_0/bias               -> alpha_net_hiku/~/equi_message_passing/~_build_diagonal/linear/b
params/message_layers_0/diagonal/layers_0/kernel             -> alpha_net_hiku/~/equi_message_passing/~_build_diagonal/linear/w
params/message_layers_0/diagonal/layers_2/bias               -> alpha_net_hiku/~/equi_message_passing/~_build_diagonal/linear_1/b
params/message_layers_0/diagonal/layers_2/kernel             -> alpha_net_hiku/~/equi_message_passing/~_build_diagonal/linear_1/w
params/message_layers_0/dir_proj/layers_0/bias               -> alpha_net_hiku/~/equi_message_passing/~_build_dir_proj/linear/b
params/message_layers_0/dir_proj/layers_0/kernel             -> alpha_net_hiku/~/equi_message_passing/~_build_dir_proj/linear/w
params/message_layers_0/dir_proj/layers_2/bias               -> alpha_net_hiku/~/equi_message_passing/~_build_dir_proj/linear_1/b
params/message_layers_0/dir_proj/layers_2/kernel             -> alpha_net_hiku/~/equi_message_passing/~_build_dir_proj/linear_1/w
params/message_layers_0/dx_layer_norm/bias                   -> alpha_net_hiku/~/equi_message_passing/layer_norm_1/offset
params/message_layers_0/dx_layer_norm/scale                  -> alpha_net_hiku/~/equi_message_passing/layer_norm_1/scale
params/message_layers_0/fc_mps/bias                          -> alpha_net_hiku/~/equi_message_passing/~/linear/b
params/message_layers_0/fc_mps/kernel                        -> alpha_net_hiku/~/equi_message_passing/~/linear/w
params/message_layers_0/kernel_imag                          -> alpha_net_hiku/~/equi_message_passing/kernel_imag
params/message_layers_0/kernel_real                          -> alpha_net_hiku/~/equi_message_passing/kernel_real
params/message_layers_0/rbf_proj/bias                        -> alpha_net_hiku/~/equi_message_passing/linear/b
params/message_layers_0/rbf_proj/kernel                      -> alpha_net_hiku/~/equi_message_passing/linear/w
params/message_layers_0/scale/bias                           -> alpha_net_hiku/~/equi_message_passing/~message/linear/b
params/message_layers_0/scale/kernel                         -> alpha_net_hiku/~/equi_message_passing/~message/linear/w
params/message_layers_0/scale2/bias                          -> alpha_net_hiku/~/equi_message_passing/linear_1/b
params/message_layers_0/scale2/kernel                        -> alpha_net_hiku/~/equi_message_passing/linear_1/w
params/message_layers_0/x_layernorm/bias                     -> alpha_net_hiku/~/equi_message_passing/layer_norm/offset
params/message_layers_0/x_layernorm/scale                    -> alpha_net_hiku/~/equi_message_passing/layer_norm/scale
params/message_layers_0/x_proj/layers_0/bias                 -> alpha_net_hiku/~/equi_message_passing/~_build_x_proj/linear/b
params/message_layers_0/x_proj/layers_0/kernel               -> alpha_net_hiku/~/equi_message_passing/~_build_x_proj/linear/w
params/message_layers_0/x_proj/layers_2/bias                 -> alpha_net_hiku/~/equi_message_passing/~_build_x_proj/linear_1/b
params/message_layers_0/x_proj/layers_2/kernel               -> alpha_net_hiku/~/equi_message_passing/~_build_x_proj/linear_1/w
params/message_layers_1/dia/bias                             -> alpha_net_hiku/~/equi_message_passing_1/~message/linear_1/b
params/message_layers_1/dia/kernel                           -> alpha_net_hiku/~/equi_message_passing_1/~message/linear_1/w
params/message_layers_1/diachi1                              -> alpha_net_hiku/~/equi_message_passing_1/diachi1
params/message_layers_1/diagonal/layers_0/bias               -> alpha_net_hiku/~/equi_message_passing_1/~_build_diagonal/linear/b
params/message_layers_1/diagonal/layers_0/kernel             -> alpha_net_hiku/~/equi_message_passing_1/~_build_diagonal/linear/w
params/message_layers_1/diagonal/layers_2/bias               -> alpha_net_hiku/~/equi_message_passing_1/~_build_diagonal/linear_1/b
params/message_layers_1/diagonal/layers_2/kernel             -> alpha_net_hiku/~/equi_message_passing_1/~_build_diagonal/linear_1/w
params/message_layers_1/dir_proj/layers_0/bias               -> alpha_net_hiku/~/equi_message_passing_1/~_build_dir_proj/linear/b
params/message_layers_1/dir_proj/layers_0/kernel             -> alpha_net_hiku/~/equi_message_passing_1/~_build_dir_proj/linear/w
params/message_layers_1/dir_proj/layers_2/bias               -> alpha_net_hiku/~/equi_message_passing_1/~_build_dir_proj/linear_1/b
params/message_layers_1/dir_proj/layers_2/kernel             -> alpha_net_hiku/~/equi_message_passing_1/~_build_dir_proj/linear_1/w
params/message_layers_1/dx_layer_norm/bias                   -> alpha_net_hiku/~/equi_message_passing_1/layer_norm_1/offset
params/message_layers_1/dx_layer_norm/scale                  -> alpha_net_hiku/~/equi_message_passing_1/layer_norm_1/scale
params/message_layers_1/fc_mps/bias                          -> alpha_net_hiku/~/equi_message_passing_1/~/linear/b
params/message_layers_1/fc_mps/kernel                        -> alpha_net_hiku/~/equi_message_passing_1/~/linear/w
params/message_layers_1/kernel_imag                          -> alpha_net_hiku/~/equi_message_passing_1/kernel_imag
params/message_layers_1/kernel_real                          -> alpha_net_hiku/~/equi_message_passing_1/kernel_real
params/message_layers_1/rbf_proj/bias                        -> alpha_net_hiku/~/equi_message_passing_1/linear/b
params/message_layers_1/rbf_proj/kernel                      -> alpha_net_hiku/~/equi_message_passing_1/linear/w
params/message_layers_1/scale/bias                           -> alpha_net_hiku/~/equi_message_passing_1/~message/linear/b
params/message_layers_1/scale/kernel                         -> alpha_net_hiku/~/equi_message_passing_1/~message/linear/w
params/message_layers_1/scale2/bias                          -> alpha_net_hiku/~/equi_message_passing_1/linear_1/b
params/message_layers_1/scale2/kernel                        -> alpha_net_hiku/~/equi_message_passing_1/linear_1/w
params/message_layers_1/x_layernorm/bias                     -> alpha_net_hiku/~/equi_message_passing_1/layer_norm/offset
params/message_layers_1/x_layernorm/scale                    -> alpha_net_hiku/~/equi_message_passing_1/layer_norm/scale
params/message_layers_1/x_proj/layers_0/bias                 -> alpha_net_hiku/~/equi_message_passing_1/~_build_x_proj/linear/b
params/message_layers_1/x_proj/layers_0/kernel               -> alpha_net_hiku/~/equi_message_passing_1/~_build_x_proj/linear/w
params/message_layers_1/x_proj/layers_2/bias                 -> alpha_net_hiku/~/equi_message_passing_1/~_build_x_proj/linear_1/b
params/message_layers_1/x_proj/layers_2/kernel               -> alpha_net_hiku/~/equi_message_passing_1/~_build_x_proj/linear_1/w
params/message_layers_2/dia/bias                             -> alpha_net_hiku/~/equi_message_passing_2/~message/linear_1/b
params/message_layers_2/dia/kernel                           -> alpha_net_hiku/~/equi_message_passing_2/~message/linear_1/w
params/message_layers_2/diachi1                              -> alpha_net_hiku/~/equi_message_passing_2/diachi1
params/message_layers_2/diagonal/layers_0/bias               -> alpha_net_hiku/~/equi_message_passing_2/~_build_diagonal/linear/b
params/message_layers_2/diagonal/layers_0/kernel             -> alpha_net_hiku/~/equi_message_passing_2/~_build_diagonal/linear/w
params/message_layers_2/diagonal/layers_2/bias               -> alpha_net_hiku/~/equi_message_passing_2/~_build_diagonal/linear_1/b
params/message_layers_2/diagonal/layers_2/kernel             -> alpha_net_hiku/~/equi_message_passing_2/~_build_diagonal/linear_1/w
params/message_layers_2/dir_proj/layers_0/bias               -> alpha_net_hiku/~/equi_message_passing_2/~_build_dir_proj/linear/b
params/message_layers_2/dir_proj/layers_0/kernel             -> alpha_net_hiku/~/equi_message_passing_2/~_build_dir_proj/linear/w
params/message_layers_2/dir_proj/layers_2/bias               -> alpha_net_hiku/~/equi_message_passing_2/~_build_dir_proj/linear_1/b
params/message_layers_2/dir_proj/layers_2/kernel             -> alpha_net_hiku/~/equi_message_passing_2/~_build_dir_proj/linear_1/w
params/message_layers_2/dx_layer_norm/bias                   -> alpha_net_hiku/~/equi_message_passing_2/layer_norm_1/offset
params/message_layers_2/dx_layer_norm/scale                  -> alpha_net_hiku/~/equi_message_passing_2/layer_norm_1/scale
params/message_layers_2/fc_mps/bias                          -> alpha_net_hiku/~/equi_message_passing_2/~/linear/b
params/message_layers_2/fc_mps/kernel                        -> alpha_net_hiku/~/equi_message_passing_2/~/linear/w
params/message_layers_2/kernel_imag                          -> alpha_net_hiku/~/equi_message_passing_2/kernel_imag
params/message_layers_2/kernel_real                          -> alpha_net_hiku/~/equi_message_passing_2/kernel_real
params/message_layers_2/rbf_proj/bias                        -> alpha_net_hiku/~/equi_message_passing_2/linear/b
params/message_layers_2/rbf_proj/kernel                      -> alpha_net_hiku/~/equi_message_passing_2/linear/w
params/message_layers_2/scale/bias                           -> alpha_net_hiku/~/equi_message_passing_2/~message/linear/b
params/message_layers_2/scale/kernel                         -> alpha_net_hiku/~/equi_message_passing_2/~message/linear/w
params/message_layers_2/scale2/bias                          -> alpha_net_hiku/~/equi_message_passing_2/linear_1/b
params/message_layers_2/scale2/kernel                        -> alpha_net_hiku/~/equi_message_passing_2/linear_1/w
params/message_layers_2/x_layernorm/bias                     -> alpha_net_hiku/~/equi_message_passing_2/layer_norm/offset
params/message_layers_2/x_layernorm/scale                    -> alpha_net_hiku/~/equi_message_passing_2/layer_norm/scale
params/message_layers_2/x_proj/layers_0/bias                 -> alpha_net_hiku/~/equi_message_passing_2/~_build_x_proj/linear/b
params/message_layers_2/x_proj/layers_0/kernel               -> alpha_net_hiku/~/equi_message_passing_2/~_build_x_proj/linear/w
params/message_layers_2/x_proj/layers_2/bias                 -> alpha_net_hiku/~/equi_message_passing_2/~_build_x_proj/linear_1/b
params/message_layers_2/x_proj/layers_2/kernel               -> alpha_net_hiku/~/equi_message_passing_2/~_build_x_proj/linear_1/w
params/message_layers_3/dia/bias                             -> alpha_net_hiku/~/equi_message_passing_3/~message/linear_1/b
params/message_layers_3/dia/kernel                           -> alpha_net_hiku/~/equi_message_passing_3/~message/linear_1/w
params/message_layers_3/diachi1                              -> alpha_net_hiku/~/equi_message_passing_3/diachi1
params/message_layers_3/diagonal/layers_0/bias               -> alpha_net_hiku/~/equi_message_passing_3/~_build_diagonal/linear/b
params/message_layers_3/diagonal/layers_0/kernel             -> alpha_net_hiku/~/equi_message_passing_3/~_build_diagonal/linear/w
params/message_layers_3/diagonal/layers_2/bias               -> alpha_net_hiku/~/equi_message_passing_3/~_build_diagonal/linear_1/b
params/message_layers_3/diagonal/layers_2/kernel             -> alpha_net_hiku/~/equi_message_passing_3/~_build_diagonal/linear_1/w
params/message_layers_3/dir_proj/layers_0/bias               -> alpha_net_hiku/~/equi_message_passing_3/~_build_dir_proj/linear/b
params/message_layers_3/dir_proj/layers_0/kernel             -> alpha_net_hiku/~/equi_message_passing_3/~_build_dir_proj/linear/w
params/message_layers_3/dir_proj/layers_2/bias               -> alpha_net_hiku/~/equi_message_passing_3/~_build_dir_proj/linear_1/b
params/message_layers_3/dir_proj/layers_2/kernel             -> alpha_net_hiku/~/equi_message_passing_3/~_build_dir_proj/linear_1/w
params/message_layers_3/dx_layer_norm/bias                   -> alpha_net_hiku/~/equi_message_passing_3/layer_norm_1/offset
params/message_layers_3/dx_layer_norm/scale                  -> alpha_net_hiku/~/equi_message_passing_3/layer_norm_1/scale
params/message_layers_3/fc_mps/bias                          -> alpha_net_hiku/~/equi_message_passing_3/~/linear/b
params/message_layers_3/fc_mps/kernel                        -> alpha_net_hiku/~/equi_message_passing_3/~/linear/w
params/message_layers_3/kernel_imag                          -> alpha_net_hiku/~/equi_message_passing_3/kernel_imag
params/message_layers_3/kernel_real                          -> alpha_net_hiku/~/equi_message_passing_3/kernel_real
params/message_layers_3/rbf_proj/bias                        -> alpha_net_hiku/~/equi_message_passing_3/linear/b
params/message_layers_3/rbf_proj/kernel                      -> alpha_net_hiku/~/equi_message_passing_3/linear/w
params/message_layers_3/scale/bias                           -> alpha_net_hiku/~/equi_message_passing_3/~message/linear/b
params/message_layers_3/scale/kernel                         -> alpha_net_hiku/~/equi_message_passing_3/~message/linear/w
params/message_layers_3/scale2/bias                          -> alpha_net_hiku/~/equi_message_passing_3/linear_1/b
params/message_layers_3/scale2/kernel                        -> alpha_net_hiku/~/equi_message_passing_3/linear_1/w
params/message_layers_3/x_layernorm/bias                     -> alpha_net_hiku/~/equi_message_passing_3/layer_norm/offset
params/message_layers_3/x_layernorm/scale                    -> alpha_net_hiku/~/equi_message_passing_3/layer_norm/scale
params/message_layers_3/x_proj/layers_0/bias                 -> alpha_net_hiku/~/equi_message_passing_3/~_build_x_proj/linear/b
params/message_layers_3/x_proj/layers_0/kernel               -> alpha_net_hiku/~/equi_message_passing_3/~_build_x_proj/linear/w
params/message_layers_3/x_proj/layers_2/bias                 -> alpha_net_hiku/~/equi_message_passing_3/~_build_x_proj/linear_1/b
params/message_layers_3/x_proj/layers_2/kernel               -> alpha_net_hiku/~/equi_message_passing_3/~_build_x_proj/linear_1/w
params/neighbor_emb/Embed_0/embedding                        -> alpha_net_hiku/~/neighbor_emb/embed/embeddings
params/radial_emb/bessel_weights                             -> alpha_net_hiku/~/rbf_emb/bessel_weights
params/radial_lin/layers_0/bias                              -> alpha_net_hiku/~/linear/b
params/radial_lin/layers_0/kernel                            -> alpha_net_hiku/~/linear/w
params/radial_lin/layers_2/bias                              -> alpha_net_hiku/~/linear_1/b
params/radial_lin/layers_2/kernel                            -> alpha_net_hiku/~/linear_1/w
params/s_vector/Dense_0/bias                                 -> alpha_net_hiku/~/s_vector/linear/b
params/s_vector/Dense_0/kernel                               -> alpha_net_hiku/~/s_vector/linear/w
params/z_emb/embedding                                       -> alpha_net_hiku/~/embed/embeddings
