如何禁用JAX在CPU环境下运行时的METAL相关提示信息？

方法1：通过环境变量强制使用CPU并屏蔽日志 #

直接在运行脚本时指定环境变量，让JAX跳过Metal后端加载并抑制警告：

JAX_PLATFORM_NAME=cpu TF_CPP_MIN_LOG_LEVEL=3 python your_script.py

如果需要在Python脚本内设置，可在导入JAX前添加：

import os
os.environ['JAX_PLATFORM_NAME'] = 'cpu'
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import jax

方法2：通过日志系统过滤特定警告 #

针对JAX的日志模块单独调整级别，屏蔽Metal相关的实验性警告：

import logging
import jax

# 屏蔽JAX平台相关的警告
logging.getLogger('jax._src.platform').setLevel(logging.ERROR)
# 屏蔽PJRT MPS客户端的警告
logging.getLogger('jax._src.pjrt').setLevel(logging.ERROR)

方法3：彻底禁用Metal后端 #

通过环境变量完全关闭JAX的Metal支持，从根源避免相关输出：

JAX_ENABLE_METAL=0 python your_script.py

或在脚本内设置：

import os
os.environ['JAX_ENABLE_METAL'] = '0'
import jax

内容的提问来源于stack exchange，提问作者João Abrantes