Fix docstring nits (#3758)

Author: ishtihoss

python/mlx/nn/layers/linear.py

@@ -115,10 +115,10 @@ class Bilinear(Module):
 
     .. math::
 
-        y_i = x_1^\top W_i x_2 + b_i
+        y_i = x_2^\top W_i x_1 + b_i
 
     where:
-    :math:`W` has shape ``[output_dims, input1_dims, input2_dims]``, :math:`b` has shape ``[output_dims ]``,
+    :math:`W` has shape ``[output_dims, input2_dims, input1_dims]``, :math:`b` has shape ``[output_dims]``,
     and :math:`i` indexes the output dimension.
 
     The values are initialized from the uniform distribution :math:`\mathcal{U}(-{k}, {k})`,

python/mlx/nn/losses.py

@@ -599,7 +599,7 @@ def margin_ranking_loss(
         >>> inputs2 = mx.array([0.75596, 0.225763, 0.256995])
         >>> loss = nn.losses.margin_ranking_loss(inputs1, inputs2, targets)
         >>> loss
-        array(0.773433, dtype=float32)
+        array([1.32937, 0.990929, 0], dtype=float32)
     """
     if not (inputs1.shape == inputs2.shape == targets.shape):
         raise ValueError(