[BUG] LayerNorm VJP returns zeros_like(weight) instead of zeros_like(bias placeholder) #3220

Describe the bug

• In LayerNorm::vjp, the df/db branch for scalar bias placeholder (b.ndim() == 0) pushes:
  • zeros_like(w, s)
• It should push:
  • zeros_like(b, s)

This creates an issue with Meta-Learning methods (MAML) or other work using Hessian vector products

To Reproduce

Include code snippet

import mlx.core as mx

D = 8
eps = 1e-5
x = mx.random.uniform(shape=(2, 4, D))
w = mx.random.uniform(shape=(D,))
y = mx.random.uniform(shape=(2, 4, D))
mx.eval(x, w, y)
  
# Second-order gradient through layer_norm with weight but no bias.
# The inner grad computes the VJP, which internally produces a bias gradient
# with shape (D,) instead of (). The outer grad differentiates through the
# VJP graph and detects the shape mismatch.

def f(x, w, y):
    return (mx.fast.layer_norm(x, w, None, eps) * y).sum()

def gf(x, w, y):
    gx, gw = mx.grad(f, argnums=(0, 1))(x, w, y)
    return ((gx + gw) * y).sum()

gx2, gw2 = mx.grad(gf, argnums=(0, 1))(x, w, y)
mx.eval(gx2, gw2)

# Crashes with:
#   ValueError: [vjp] Output shape (8) does not match cotangent shape ().
#
# The reference (non-fast) layer_norm handles this correctly.

Expected behavior
Should be handled in the same way as the non-fast layer_norm...

Desktop (please complete the following information):

• OS Version: 26.3.1

Additional context

• Note: in typical Python usage (bias=None), the framework does not request argnum 2,
  so the buggy vjps[2] is computed but discarded. The bug surfaces when the bias gradient
  is explicitly requested or under transform composition.
• This is NOT a synthetic edge case. Second-order gradients through layer_norm are
  used in meta-learning (MAML), physics-informed neural networks, and Hessian
  computation. Any model that uses layer_norm with weight but no bias and needs
  second-order differentiation will hit this crash.
• The reference implementation (pure Python layer_norm) handles the same second-order
  gradient correctly, confirming this is a bug in the fast path, not a limitation.