Jax transforms and modules cannot be mixed.

[yang-song]

It sucks to know this after investing so much time in Flax hoping to start a new project. I just wish it was possible to use jax.lax.cond inside a module where branches can contain flax modules as well. It came as a big surprise to me that this was not supported

[marcvanzee]

Hi @yang-song I'm sorry to hear that! Do you have a code example? I'm curious how you were planning to use the module in jax.lax.cond.

[yang-song]

Thanks for the fast response! The following code is part of the effort on implementing PixelCNN++ with caching (for faster sampling). The module may or may not compute the convolution, depending on the pixel location. WNConv is a convolutional layer defined by myself that needs to appear inside jax.lax.cond

class down_shifted_conv2d(nn.Module):
  num_filters_out: int
  batch_size: int
  image_width: int
  filter_size: Tuple[int, int] = (2, 3)
  stride: Tuple[int, int] = (1, 1)

  @nn.compact
  def __call__(self, row_input, row, col, cache_every, run_every):
    filter_height = self.filter_size[0]
    filter_width = self.filter_size[1]
    cache_height = filter_height
    padding = filter_width // 2
    cache_width = self.image_width + 2 * padding
    output_width = int(math.ceil(self.image_width / float(self.stride[0])))

    cache = self.variable('cache', 'cache', _init_cache,
                          self.batch_size, cache_height,
                          cache_width, row_input.shape[-1])

    output_cache = self.variable('cache', 'output_cache', _init_cache,
                                 self.batch_size, 1, output_width,
                                 self.num_filters_out)

    cache.value = jax.lax.cond(row == 0 and col == 0,
                               lambda cache: jnp.zeros_like(cache),
                               lambda cache: cache, cache.value)

    def col_0(caches):
      cache, output_cache = caches
      cache = jax.lax.cond(row % cache_every == 0,
                           lambda c: c.at[:, -1:, padding:(padding + self.image_width), :].set(row_input),
                           lambda c: c, cache)

      def row_run_every(caches):
        cache, output_cache = caches
        outputs = WNConv(features=self.num_filters_out,
                         kernel_size=self.filter_size,
                         strides=self.stride)(cache)
        return outputs

      output_cache = jax.lax.cond(row % run_every == 0,
                                  row_run_every, lambda x: x[1],
                                  (cache, output_cache))

      cache = jax.lax.cond(row % cache_every == 0,
                           lambda c: _roll_cache(c),
                           lambda c: c,
                           cache)

      return (cache, output_cache)

    (cache.value, output_cache.value) = jax.lax.cond(col == 0,
                                                     col_0,
                                                     lambda x: x,
                                                     (cache.value, output_cache.value))
    return output_cache.value

[jheek]

We will add full fledged cond support that would allow module-in-cond.

For your particular case though I would advice not to use cond at all.
Cond has a high overhead on accelerators in particular GPU this is also why it hasn't been implemented yet. It is rarely a good idea to use it inside a model. The condition needs to be evaluated on host and the next op can only be scheduled afterward causing the GPU to idle in the meantime.

jnp.where will get you better performance here. It will execute both branches and pick the correct one avoiding a sync between GPU and host while also allowing for more optimization in XLA.

Alternatively, you could use a boolean first_row, first_col and use jit(static_argnums=()) to compile the 4 cases of the sampling code. This might be even faster than select and will avoid the idle GPU problem as well.

[jheek]

If you really want to use cond right away you can use the following trick though:

# inside @compact 

conv_params = self.param('cond_conv', partial(WNConv.init, **hparams), x)
lax.cond(condition, lambda x: WNConv.apply(conv_params, x, **hparams)), ...)

[yang-song]

Thank you for such a great answer!