import jax
import jax.numpy as jnp
from jax import random
import jax.nn.initializers as init

# 设置随机数种子
key = random.PRNGKey(42)

# 定义权重的形状
shape = (in_dim, out_dim)

# 获取 Glorot 正态初始化函数
glorot_normal_init = init.glorot_normal()

# 初始化权重
weights = glorot_normal_init(key, shape)

指定权重矩阵shape.

使用init.glorot_normal()获取Glorot正态初始化函数.

Glorot正态初始化的方差由输入和输出的神经元数量决定：

jax.nn.initializers.glorot_normal()是JAX库中的一个函数，用于初始化神经网络的参数。它使用Glorot正态分布初始化方法，也称为Xavier正态初始化。这种初始化方法旨在使每个神经元的输出具有相同的方差，以促进梯度在网络中流动时的稳定性。