Function like keras.layers.Masking

[dev-sora]

Hi,

Is there any function like keras.layers.Masking, which masks padded value in input sequences?
Or can we make custom layer working same? If yes, I'm really appreciate if you tell me how.

[matthias-wright]

Hi @dev-sora, Flax doesn't have a function like that.
Could you give an example of your specific use case? Is it for RNN layers?

[dev-sora]

Hi @matthias-wright , thanks for your responding.
I wanna use a function like padding to input data acquired using a sensor with an unstable sampling rate into a neural network.

[dev-sora]

latter.
I believe it is relatively common to use sensors with unstable sampling rates, but in such cases, they use linear interpolation or other methods to align the length of the data. I wanna use padding because the data I am dealing with cannot be linearly interpolated for some reason, but I can only find examples of padding/masking use in the context of natural language processing.

[matthias-wright]

I see. I think that tf.keras.layers.Masking in intended for RNN layers. Since RNN layers process the input sequentially, they can simply skip masked values at a specific timestep. However, when you process your data all at once with fc layers, you cannot easily ignore certain values in the input.

[dev-sora]

Is it impossible to merge keras function into NN model written in Flax just because keras is now one part of tensorflow?

[matthias-wright]

You can't simply plug Keras layers into a Flax model if you mean that. You can have a look here for interoperation between TF and Jax. I have never experimented with that though.

[dev-sora]

Thank you. I'll investigate more about it.

[jcrichard]

I was asking myself the same question. Thanks, @matthias-wright for pointing out jax2tf.

By the way, tf.keras.layers.Masking is only used when a mask is in the Call argument of the next layer as it is the case for LSTM.

[marcvanzee]

@dev-sora, I think you are saying you want to mask out a batch of sequences after a certain length, where these lengths can differ per example, is that correct? Will this work for you?

def mask_sequences(sequences, *, eos_id=0):
  lengths = jnp.where(sequences == eos_id)[-1]
  lengths = jnp.expand_dims(lengths, -1)
  max_len = jnp.arange(sequences.shape[1])
  return sequences * (lengths > max_len)
  
sequences = jnp.array([[1, 2, 3, 0, 5, 6],
                       [1, 0, 3, 4, 5, 6]])

mask_sequences(sequences)

Returns

DeviceArray([[1, 2, 3, 0, 0, 0],
             [1, 0, 0, 0, 0, 0]], dtype=int32)

This will mask everything after the occurrence of eos_id out with 0. If you want to use a different value than 0 you could follow this up by replacing all occurrences of 0 with your desired value using jnp.where.