Make shape a tuple (#591)

* shape tuple

* also remove simplify from docs

* rebase

python/mlx/nn/layers/dropout.py

@@ -81,7 +81,7 @@ class Dropout2d(Module):
         # Dropout is applied on the whole channel
         # 3D input: (1, 1, C)
         # 4D input: (B, 1, 1, C)
-        mask_shape = x.shape
+        mask_shape = list(x.shape)
         mask_shape[-2] = mask_shape[-3] = 1
 
         mask = mx.random.bernoulli(p=self._p_1, shape=mask_shape)

python/mlx/nn/losses.py

@@ -70,8 +70,9 @@ def cross_entropy(
     targets_as_probs = targets.ndim == logits.ndim
 
     def _drop_dim(shape, axis):
+        shape = list(shape)
         shape.pop(axis)
-        return shape
+        return tuple(shape)
 
     # Check shapes in two cases: targets as class indices and targets as probabilities
     if (targets_as_probs and targets.shape != logits.shape) or (