Bug fix in conjugate_gradient
Performing multiple iterations on a zero residual now avoids NaN.
This problem could arise when batch_size > 0 and one example has k=0.
@@ -45,6 +45,7 @@ def conjugate_gradient(k, apply_A, initial_x=None, accuracy=1e-5, max_iterations
residual : residual
"""
tmp = math.sum(momentum * A_times_momentum, axis=1, keepdims=True) # t = sum(mAm)
+ tmp = math.where(math.equal(tmp, 0), math.ones_like(tmp), tmp)
a = math.sum(momentum * residual, axis=1, keepdims=True) / tmp # a = sum(mr)/sum(mAm)
pressure += a * momentum # p += am
residual -= a * A_times_momentum # r -= aAm
