def backward_pass(X, y, z1, a1, z2, a2, W1, b1, W2, b2):
    """Compute gradients using backpropagation"""
    m = X.shape[0]  # Number of samples

    # Output layer gradients
    dz2 = 2 * (a2 - y) * sigmoid_derivative(z2)  # (4, 1)
    dW2 = a1.T @ dz2 / m                         # (4, 1)
    db2 = np.mean(dz2, axis=0, keepdims=True)    # (1, 1)

    # Hidden layer gradients
    dz1 = (dz2 @ W2.T) * sigmoid_derivative(z1)  # (4, 4)
    dW1 = X.T @ dz1 / m                          # (2, 4)
    db1 = np.mean(dz1, axis=0, keepdims=True)    # (1, 4)

    return dW1, db1, dW2, db2

# Test backpropagation
dW1, db1, dW2, db2 = backward_pass(X, y, z1, a1, z2, predictions, W1, b1, W2, b2)
print(f"Gradient shapes:")
print(f"  dW1: {dW1.shape}, dW2: {dW2.shape}")
print(f"  db1: {db1.shape}, db2: {db2.shape}")