jax.nn.get_scaled_dot_general_config

jax.nn.get_scaled_dot_general_config(mode, global_scale=None)

获取scaled_dot_general的量化配置。

为jax.nn.scaled_dot_general创建量化配置。

另请参阅

• jax.nn.scaled_dot_general(): 缩放点积通用函数。

参数:

• mode (Literal['nvfp4', 'mxfp8'])

• global_scale (Array | None)