I am fairly new to optax. There is this line in scale_by_lbfgs that does the weight calculation in optax/optax/_src/transform.py that's causing some floating point errors for me: def scale_by_lbfgs( ... def update_fn( ... weight = jnp.where( vdot_diff_params_updates == 0.0, 0.0, 1.0 / vdot_diff_params_updates ) This gives nan when vdot_diff_params_updates is small but non-zero. Is there a good way to avoid this, such as editing states or gradient capping? Continue reading...
