TangTangFei

i'm getting similar results: for example for the test `test_mamba_inner_fn` with `is_variable_B=True, is_variable_C=True, seqlen=128, itype=torch.float32, wtype=torch.complex64` ``` Output max diff: 0.0 Output mean diff: 0.0 dxz max diff: 5888.0 dx...

i have changed the scale of the following tensors: ``` xz = (0.01 * torch.rand(bs, 2 * dim, seq_len, dtype=itype, device=device)).requires_grad_() g = torch.randn_like(out) * 0.01 ``` the discrepancies become...