Fixing off-diag impl, adding test case.