How to bypass constructing arrays via boolean masking

[qsh-zh]

The question is closely related to the topic. I am trying to do some calculation that depends on other input conditions. Here is a toy example

## without jit
def loss_fn(param, x):
  mask = some_cond_fn(x) # array of boolean
  real_input = x[mask]
  return expensive_loss_fn(param, x).sum()

However, if I want to use jax.jit, I can not construct real_input via boolean masking. There is another way to bypass the issue but pay more costs

def loss_fn(param, x):
  losses = expensive_loss_fn(param, x) # some calculation is waste here
  mask = some_cond_fn(x) # array of boolean
  return (mask * losses).sum()

Do we have an efficient way to do it?

[jakevdp]

Unfortunately, the answer will depend on the details of expensive_loss_function. Can you provide a more complete example of the kind of computation you're hoping to do?

[qsh-zh]

@jakevdp Thanks for your help. I am using a neural network based on flax for expensive_loss_function, you can think it is a large CNN or MLP, minimizing some KL divergence or quadratic loss.

[jakevdp]

OK, thanks. There is currently no general way to do what you want (i.e. create a dynamically-shaped array within a JIT context and pass it to a black-box JIT-compiled function). The solution you came up with (computing the result and then applying the mask to the sum) is probably the best you can do.

Alternatively, if you are able to reach into the implementation of your black box function, you can use other tricks for making the computation more efficient given masked input. But unfortunately, it's hard to say much in general about that, because the possible approaches depend tightly on what computations are being performed.

Hope that helps

[qsh-zh]

I see. I want to skip some data in one batch. Do you have any ideas on how to do the task, for example, like CNN in flax tutorial

class CNN(nn.Module):
  """A simple CNN model."""

  @nn.compact
  def __call__(self, x):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    return x


@jax.jit
def apply_model(state, images, labels):
  """Computes gradients, loss and accuracy for a single batch."""
  def loss_fn(params):
    logits = CNN().apply({'params': params}, images)
    one_hot = jax.nn.one_hot(labels, 10)
    loss = jnp.mean(optax.softmax_cross_entropy(logits=logits, labels=one_hot))
    return loss, logits

  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
  (loss, logits), grads = grad_fn(state.params)
  accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
  return grads, loss, accuracy

[jakevdp]

I don't know of any easy way to make a flax model handle a dynamically-shaped input when called within a JIT context.

[YouJiacheng]

You can first sort x according to mask, then use lax.while_loop to iterate over valid part. With blocked loop(i.e. process >1 sample per iteration), you can tradeoff between parallelism and computation waste.
However, lax.while_loop doesn't support reverse-mode autodiff, you might need custom_vjp.

[qsh-zh]

@YouJiacheng , thanks for your suggestions. But I feel the approach lost parallel calculation of GPU if I use large networks and maybe slower.