[BUG] Dropout not preserving `dtype` about mlx

[bhargavyagnik]

I experimented two approaches for

mlx/python/mlx/nn/layers/dropout.py

Line 35 in 1086dc4

return (1 / self._p_1) * mask * x

• return (1 / self._p_1) * mask * x # Return dtype is mx.float32

but when I replace that with

• return mask * x * ( 1 / self._p_1) # Return dtype is mx.bfloat16

This problem is related to type promotion and similar usecase exist in Numpy and numerous libraries.

Since python executes from left to right; So in the first case, when weakly typed float (1 / self._p_1) is multiplied with mx.bool_, (mask) the resulting dtype is mx.float32 and thus entire equation is calculated in mx.float32.

When we first multiply mx.bool_ with mx.bfloat16, the weakly typed float is promoted to just mx.bfloat16thus preserving the dtype

[vgoklani]

class Dropout(nn.Module):
    def __init__(self, p: Optional[float] = 0.5):
        super().__init__()
        self.p = p

    def __call__(self, x: mx.array, key: Optional[mx.array] = None):
        if not self.training or self.p == 0:
            return x
        else:
            mask = mx.random.uniform(shape=x.shape, dtype=x.dtype, key=key) > self.p
            return mask * x / (1.0 - self.p)

[vgoklani]

Just a heads up, the tree_flaten method drops the Dropout layers from the network, so when we serialized the optimizer:

mx.save_safetensors(file=optimizer_state_path, arrays=dict(optimizer_state))

and decided to continue training from a previously saved checkpoint, we got a bunch of key errors...

[awni]

I'm not following. The Dropout layers shouldn't be saved in tree_flatten since they don't have parameters. Could you share some code which reproduces the errors you are getting?

[vgoklani]

@awni created a separate issue for this: #1328