I’m looking to mask the input to a sequence model with independently randomly placed blocks of random length.  Here is a function prototype with pseudo-code for what I want:
def mask(x, prob, max_length, batch_dim=0, seq_dim=1, mask_value=0):
    """Returns a new tensor like x but possibly with certain elements masked to mask_value.

    More precisely,
      ret = x.clone()
      for each i < x.size(batch_dim)
        for each j < x.size(seq_dim)
          with probability prob
            l = randint(1, max_length)
            for each index tuple idx of ret
              if idx[batch_dim] == i and j <= idx[seq_dim] < j+l
                ret[idx] = mask_value
      return ret
    """
    pass

I am looking for a way to do this without actually looping through the tensor, which is at least inelegant and possibly would cause a speed issue for the model (didn’t profile it though).
I’m also open to suggestions of variant masking strategies that might work equally well but be easier to implement if this one is hard.