theano Loops with theano making while loop

Example

As of theano 0.9, while loops can be done via theano.scan_module.scan_utils.until. To use, you should return until object in fn of scan.

In the following example, we build a function that checks whether a complex number is inside Mandelbrot set. A complex number z_0 is inside mandelbrot set if series z_{n+1} = z_{n}^2 + z_0 does not converge.

MAX_ITER = 256
BAILOUT = 2.
s_z0 = th.cscalar()
def iterate(s_i_, s_z_, s_z0_):
    return [s_z_*s_z_+s_z0_,s_i_+1], {}, until(T.abs_(s_z_)>BAILOUT)
(_1, s_niter), _2 = theano.scan(
    fn = iterate,
    outputs_info = [0, s_z0],
    non_sequences = [s_z0],
    n_steps = MAX_ITER
)
fn_mandelbrot_iters = theano.function([s_z0], s_niter)
def is_in_mandelbrot(z_):
    return fn_mandelbrot_iters(z_)>=MAX_ITER

is_in_mandelbrot(0.24+0.j) # True
is_in_mandelbrot(1.j) # True
is_in_mandelbrot(0.26+0.j) # False