JAX类方法中使用scan_tqdm进度条的问题及解决方案

问题描述 #

在使用JAX框架时，希望为类方法添加进度条。目标方法定义如下：

@partial(jit, static_argnums=(0,))
def _run_step(self, runner, i_step):
           ...

尝试使用scan_tqdm装饰器后未生效，原因是该装饰器要求被装饰函数仅接收两个输入：函数业务输入（此处为runner）和进度条步骤计数器（此处为i_step），但类方法自带的self参数不符合这个要求，导致scan_tqdm无法正常工作。

解决办法 #

最终通过在lax.scan调用中直接嵌套scan_tqdm的方式解决了问题，代码示例如下：

runner, metrics = lax.scan(
                scan_tqdm(self.config["TOTAL_STEPS"])(self._run_step),
                runner,
                jnp.arange(self.config["TOTAL_STEPS"]),
                self.config["TOTAL_STEPS"]
            )

内容的提问来源于stack exchange，提问作者amavrits