PyTorchFlaxӳ:
--------------------------------------------------------------------------------
model.a -> a
model.b -> b
model.kernel1 -> kernel1
model.kernels_real -> kernels_real
model.kernels_imag -> kernels_imag
model.z_emb.weight -> z_emb/embedding
model.radial_emb.bessel_weights -> radial_emb/bessel_weights
model.radial_lin.0.weight -> radial_lin/layers_0/kernel
model.radial_lin.0.bias -> radial_lin/layers_0/bias
model.radial_lin.2.weight -> radial_lin/layers_2/kernel
model.radial_lin.2.bias -> radial_lin/layers_2/bias
model.neighbor_emb.embedding.weight -> neighbor_emb/Embed_0/embedding
model.S_vector.lin1.0.weight -> s_vector/Dense_0/kernel
model.S_vector.lin1.0.bias -> s_vector/Dense_0/bias
model.lin.0.weight -> lin/layers_0/kernel
model.lin.0.bias -> lin/layers_0/bias
model.lin.2.weight -> lin/layers_2/kernel
model.lin.2.bias -> lin/layers_2/bias
model.message_layers.0.diachi1 -> message_layers_0/diachi1
model.message_layers.0.kernel_real -> message_layers_0/kernel_real
model.message_layers.0.kernel_imag -> message_layers_0/kernel_imag
model.message_layers.0.unitary -> message_layers_0/unitary
model.message_layers.0.scale.weight -> message_layers_0/scale/kernel
model.message_layers.0.scale.bias -> message_layers_0/scale/bias
model.message_layers.0.dir_proj.0.weight -> message_layers_0/dir_proj/layers_0/kernel
model.message_layers.0.dir_proj.0.bias -> message_layers_0/dir_proj/layers_0/bias
model.message_layers.0.dir_proj.2.weight -> message_layers_0/dir_proj/layers_2/kernel
model.message_layers.0.dir_proj.2.bias -> message_layers_0/dir_proj/layers_2/bias
model.message_layers.0.x_proj.0.weight -> message_layers_0/x_proj/layers_0/kernel
model.message_layers.0.x_proj.0.bias -> message_layers_0/x_proj/layers_0/bias
model.message_layers.0.x_proj.2.weight -> message_layers_0/x_proj/layers_2/kernel
model.message_layers.0.x_proj.2.bias -> message_layers_0/x_proj/layers_2/bias
model.message_layers.0.rbf_proj.weight -> message_layers_0/rbf_proj/kernel
model.message_layers.0.rbf_proj.bias -> message_layers_0/rbf_proj/bias
model.message_layers.0.x_layernorm.weight -> message_layers_0/x_layernorm/scale
model.message_layers.0.x_layernorm.bias -> message_layers_0/x_layernorm/bias
model.message_layers.0.diagonal.0.weight -> message_layers_0/diagonal/layers_0/kernel
model.message_layers.0.diagonal.0.bias -> message_layers_0/diagonal/layers_0/bias
model.message_layers.0.diagonal.2.weight -> message_layers_0/diagonal/layers_2/kernel
model.message_layers.0.diagonal.2.bias -> message_layers_0/diagonal/layers_2/bias
model.message_layers.0.dx_layer_norm.weight -> message_layers_0/dx_layer_norm/scale
model.message_layers.0.dx_layer_norm.bias -> message_layers_0/dx_layer_norm/bias
model.message_layers.0.scale2.0.weight -> message_layers_0/scale2/kernel
model.message_layers.0.scale2.0.bias -> message_layers_0/scale2/bias
model.message_layers.0.fc_mps.weight -> message_layers_0/fc_mps/kernel
model.message_layers.0.fc_mps.bias -> message_layers_0/fc_mps/bias
model.message_layers.0.fc_dx.weight -> message_layers_0/fc_dx/kernel
model.message_layers.0.fc_dx.bias -> message_layers_0/fc_dx/bias
model.message_layers.0.dia.weight -> message_layers_0/dia/kernel
model.message_layers.0.dia.bias -> message_layers_0/dia/bias
model.message_layers.1.diachi1 -> message_layers_1/diachi1
model.message_layers.1.kernel_real -> message_layers_1/kernel_real
model.message_layers.1.kernel_imag -> message_layers_1/kernel_imag
model.message_layers.1.unitary -> message_layers_1/unitary
model.message_layers.1.scale.weight -> message_layers_1/scale/kernel
model.message_layers.1.scale.bias -> message_layers_1/scale/bias
model.message_layers.1.dir_proj.0.weight -> message_layers_1/dir_proj/layers_0/kernel
model.message_layers.1.dir_proj.0.bias -> message_layers_1/dir_proj/layers_0/bias
model.message_layers.1.dir_proj.2.weight -> message_layers_1/dir_proj/layers_2/kernel
model.message_layers.1.dir_proj.2.bias -> message_layers_1/dir_proj/layers_2/bias
model.message_layers.1.x_proj.0.weight -> message_layers_1/x_proj/layers_0/kernel
model.message_layers.1.x_proj.0.bias -> message_layers_1/x_proj/layers_0/bias
model.message_layers.1.x_proj.2.weight -> message_layers_1/x_proj/layers_2/kernel
model.message_layers.1.x_proj.2.bias -> message_layers_1/x_proj/layers_2/bias
model.message_layers.1.rbf_proj.weight -> message_layers_1/rbf_proj/kernel
model.message_layers.1.rbf_proj.bias -> message_layers_1/rbf_proj/bias
model.message_layers.1.x_layernorm.weight -> message_layers_1/x_layernorm/scale
model.message_layers.1.x_layernorm.bias -> message_layers_1/x_layernorm/bias
model.message_layers.1.diagonal.0.weight -> message_layers_1/diagonal/layers_0/kernel
model.message_layers.1.diagonal.0.bias -> message_layers_1/diagonal/layers_0/bias
model.message_layers.1.diagonal.2.weight -> message_layers_1/diagonal/layers_2/kernel
model.message_layers.1.diagonal.2.bias -> message_layers_1/diagonal/layers_2/bias
model.message_layers.1.dx_layer_norm.weight -> message_layers_1/dx_layer_norm/scale
model.message_layers.1.dx_layer_norm.bias -> message_layers_1/dx_layer_norm/bias
model.message_layers.1.scale2.0.weight -> message_layers_1/scale2/kernel
model.message_layers.1.scale2.0.bias -> message_layers_1/scale2/bias
model.message_layers.1.fc_mps.weight -> message_layers_1/fc_mps/kernel
model.message_layers.1.fc_mps.bias -> message_layers_1/fc_mps/bias
model.message_layers.1.fc_dx.weight -> message_layers_1/fc_dx/kernel
model.message_layers.1.fc_dx.bias -> message_layers_1/fc_dx/bias
model.message_layers.1.dia.weight -> message_layers_1/dia/kernel
model.message_layers.1.dia.bias -> message_layers_1/dia/bias
model.message_layers.2.diachi1 -> message_layers_2/diachi1
model.message_layers.2.kernel_real -> message_layers_2/kernel_real
model.message_layers.2.kernel_imag -> message_layers_2/kernel_imag
model.message_layers.2.unitary -> message_layers_2/unitary
model.message_layers.2.scale.weight -> message_layers_2/scale/kernel
model.message_layers.2.scale.bias -> message_layers_2/scale/bias
model.message_layers.2.dir_proj.0.weight -> message_layers_2/dir_proj/layers_0/kernel
model.message_layers.2.dir_proj.0.bias -> message_layers_2/dir_proj/layers_0/bias
model.message_layers.2.dir_proj.2.weight -> message_layers_2/dir_proj/layers_2/kernel
model.message_layers.2.dir_proj.2.bias -> message_layers_2/dir_proj/layers_2/bias
model.message_layers.2.x_proj.0.weight -> message_layers_2/x_proj/layers_0/kernel
model.message_layers.2.x_proj.0.bias -> message_layers_2/x_proj/layers_0/bias
model.message_layers.2.x_proj.2.weight -> message_layers_2/x_proj/layers_2/kernel
model.message_layers.2.x_proj.2.bias -> message_layers_2/x_proj/layers_2/bias
model.message_layers.2.rbf_proj.weight -> message_layers_2/rbf_proj/kernel
model.message_layers.2.rbf_proj.bias -> message_layers_2/rbf_proj/bias
model.message_layers.2.x_layernorm.weight -> message_layers_2/x_layernorm/scale
model.message_layers.2.x_layernorm.bias -> message_layers_2/x_layernorm/bias
model.message_layers.2.diagonal.0.weight -> message_layers_2/diagonal/layers_0/kernel
model.message_layers.2.diagonal.0.bias -> message_layers_2/diagonal/layers_0/bias
model.message_layers.2.diagonal.2.weight -> message_layers_2/diagonal/layers_2/kernel
model.message_layers.2.diagonal.2.bias -> message_layers_2/diagonal/layers_2/bias
model.message_layers.2.dx_layer_norm.weight -> message_layers_2/dx_layer_norm/scale
model.message_layers.2.dx_layer_norm.bias -> message_layers_2/dx_layer_norm/bias
model.message_layers.2.scale2.0.weight -> message_layers_2/scale2/kernel
model.message_layers.2.scale2.0.bias -> message_layers_2/scale2/bias
model.message_layers.2.fc_mps.weight -> message_layers_2/fc_mps/kernel
model.message_layers.2.fc_mps.bias -> message_layers_2/fc_mps/bias
model.message_layers.2.fc_dx.weight -> message_layers_2/fc_dx/kernel
model.message_layers.2.fc_dx.bias -> message_layers_2/fc_dx/bias
model.message_layers.2.dia.weight -> message_layers_2/dia/kernel
model.message_layers.2.dia.bias -> message_layers_2/dia/bias
model.message_layers.3.diachi1 -> message_layers_3/diachi1
model.message_layers.3.kernel_real -> message_layers_3/kernel_real
model.message_layers.3.kernel_imag -> message_layers_3/kernel_imag
model.message_layers.3.unitary -> message_layers_3/unitary
model.message_layers.3.scale.weight -> message_layers_3/scale/kernel
model.message_layers.3.scale.bias -> message_layers_3/scale/bias
model.message_layers.3.dir_proj.0.weight -> message_layers_3/dir_proj/layers_0/kernel
model.message_layers.3.dir_proj.0.bias -> message_layers_3/dir_proj/layers_0/bias
model.message_layers.3.dir_proj.2.weight -> message_layers_3/dir_proj/layers_2/kernel
model.message_layers.3.dir_proj.2.bias -> message_layers_3/dir_proj/layers_2/bias
model.message_layers.3.x_proj.0.weight -> message_layers_3/x_proj/layers_0/kernel
model.message_layers.3.x_proj.0.bias -> message_layers_3/x_proj/layers_0/bias
model.message_layers.3.x_proj.2.weight -> message_layers_3/x_proj/layers_2/kernel
model.message_layers.3.x_proj.2.bias -> message_layers_3/x_proj/layers_2/bias
model.message_layers.3.rbf_proj.weight -> message_layers_3/rbf_proj/kernel
model.message_layers.3.rbf_proj.bias -> message_layers_3/rbf_proj/bias
model.message_layers.3.x_layernorm.weight -> message_layers_3/x_layernorm/scale
model.message_layers.3.x_layernorm.bias -> message_layers_3/x_layernorm/bias
model.message_layers.3.diagonal.0.weight -> message_layers_3/diagonal/layers_0/kernel
model.message_layers.3.diagonal.0.bias -> message_layers_3/diagonal/layers_0/bias
model.message_layers.3.diagonal.2.weight -> message_layers_3/diagonal/layers_2/kernel
model.message_layers.3.diagonal.2.bias -> message_layers_3/diagonal/layers_2/bias
model.message_layers.3.dx_layer_norm.weight -> message_layers_3/dx_layer_norm/scale
model.message_layers.3.dx_layer_norm.bias -> message_layers_3/dx_layer_norm/bias
model.message_layers.3.scale2.0.weight -> message_layers_3/scale2/kernel
model.message_layers.3.scale2.0.bias -> message_layers_3/scale2/bias
model.message_layers.3.fc_mps.weight -> message_layers_3/fc_mps/kernel
model.message_layers.3.fc_mps.bias -> message_layers_3/fc_mps/bias
model.message_layers.3.fc_dx.weight -> message_layers_3/fc_dx/kernel
model.message_layers.3.fc_dx.bias -> message_layers_3/fc_dx/bias
model.message_layers.3.dia.weight -> message_layers_3/dia/kernel
model.message_layers.3.dia.bias -> message_layers_3/dia/bias
model.FTEs.0.vec_proj.weight -> ftes_0/vec_proj/kernel
model.FTEs.0.xvec_proj.0.weight -> ftes_0/xvec_proj/layers_0/kernel
model.FTEs.0.xvec_proj.0.bias -> ftes_0/xvec_proj/layers_0/bias
model.FTEs.0.xvec_proj.2.weight -> ftes_0/xvec_proj/layers_2/kernel
model.FTEs.0.xvec_proj.2.bias -> ftes_0/xvec_proj/layers_2/bias
model.FTEs.1.vec_proj.weight -> ftes_1/vec_proj/kernel
model.FTEs.1.xvec_proj.0.weight -> ftes_1/xvec_proj/layers_0/kernel
model.FTEs.1.xvec_proj.0.bias -> ftes_1/xvec_proj/layers_0/bias
model.FTEs.1.xvec_proj.2.weight -> ftes_1/xvec_proj/layers_2/kernel
model.FTEs.1.xvec_proj.2.bias -> ftes_1/xvec_proj/layers_2/bias
model.FTEs.2.vec_proj.weight -> ftes_2/vec_proj/kernel
model.FTEs.2.xvec_proj.0.weight -> ftes_2/xvec_proj/layers_0/kernel
model.FTEs.2.xvec_proj.0.bias -> ftes_2/xvec_proj/layers_0/bias
model.FTEs.2.xvec_proj.2.weight -> ftes_2/xvec_proj/layers_2/kernel
model.FTEs.2.xvec_proj.2.bias -> ftes_2/xvec_proj/layers_2/bias
model.FTEs.3.vec_proj.weight -> ftes_3/vec_proj/kernel
model.FTEs.3.xvec_proj.0.weight -> ftes_3/xvec_proj/layers_0/kernel
model.FTEs.3.xvec_proj.0.bias -> ftes_3/xvec_proj/layers_0/bias
model.FTEs.3.xvec_proj.2.weight -> ftes_3/xvec_proj/layers_2/kernel
model.FTEs.3.xvec_proj.2.bias -> ftes_3/xvec_proj/layers_2/bias
model.last_layer.weight -> last_layer/kernel
model.last_layer.bias -> last_layer/bias
model.last_layer_quantum.weight -> last_layer_quantum/kernel
model.last_layer_quantum.bias -> last_layer_quantum/bias
