As usual, Sayak writes tutorials faster than I can read them .
This is not the implementation I would expect for stocastic-depth. Isn’t this more more like … dropout with a funny noise-shape . I expected:
class Stocastic(keras.Layer):
  def __init__(self, wrapped, drop_prob, *args, **kwargs):
    super().__init(*args, **kwargs)
    self.wrapped = wrapped
    self.drop_prob = drop_prob

  def call(self, x):
    if tf.random.uniform(shape=()) > self.drop_prob:
      x = self.wrapped(x)

    return x

Otherwise you don’t get all the training speed improvements they talk about in the stocastic-depth paper since you run the layer anyway. Also the dropout-like implementation… this only works to kill the output before the add on a residual branch, right? Does the dropout 1/keep_prob scaling still make sense used like this?