RMSNorm Backward Pass

Problem

Given upstream derivative dy for y = x * weight / sqrt(mean(x²)+eps), return {"dx": gradient with respect to x, "dweight": gradient with respect to weight}. Sum the shared weight gradient across all rows. Do not numerically differentiate in your submission.

The derivative with respect to x includes the dependence of the denominator on every feature in its row.

def rms_norm_backward(x, weight, dy, eps):

Constraints

• x and dy are N × D; weight has D entries; 1 ≤ N,D ≤ 8.
• eps ≥ 0 and every row denominator is positive.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    [
      2
    ]
  ],
  [
    1
  ],
  [
    [
      1
    ]
  ],
  0
]
Expected result: {
  "dx": [
    [
      0
    ]
  ],
  "dweight": [
    1.000000000001
  ]
}

Arguments: [
  [
    [
      1,
      0
    ]
  ],
  [
    1,
    1
  ],
  [
    [
      0,
      1
    ]
  ],
  0
]
Expected result: {
  "dx": [
    [
      0,
      1.4142135623023844
    ]
  ],
  "dweight": [
    0,
    0
  ]
}