在Tensorflow中做Mandlebrot示例的更有效方法

Question

The tensorflow website has an example where TF is used to compute the Mandlebrot set. Here's the relevant snippet:

# Compute the new values of z: z^2 + x
zs_ = zs*zs + xs

# Have we diverged with this new value?
not_diverged = tf.abs(zs_) < 4

# Operation to update the zs and the iteration count.
#
# Note: We keep computing zs after they diverge! This
#       is very wasteful! There are better, if a little
#       less simple, ways to do this.
#
step = tf.group(
  zs.assign(zs_),
  ns.assign_add(tf.cast(not_diverged, tf.float32))
  )

The relevant quote is "There are better, if a little less simple, ways to do this." Does anyone know what a better way might be?

I'm messing around with ray tracing in TF, and I'm encountering situations where 90% of pixels have converged, but I keep recomputing them because I don't know how to update a subset of the entires in a tensor without sacrificing the speed benefits of using vector operations everywhere.

Answer 1

I figured out a way of doing it.

The key is to use tf.where to find the pixels that haven't yet diverged, tf.gather_nd to pull those pixels into a smaller array, perform the update step on those specific pixels, then use tf.scatter_nd (or another "scatter" function) to apply sparse updates to some variable representing the state.

In my use case, this saved quite a lot of time.