def METHOD_NAME(self):
    def check(a):
        perms = list(permutations(range(a.ndim))) + [None]
        for perm in perms:
            x = jt.array(a).float()
            if perm:
                y = x.transpose(perm)
            else:
                y = x.transpose()
            dx = jt.grad(y*y, x).data
            self.assertEqual(dx.shape, a.shape)
            assert (dx==a*2).all(), f"\n{dx}\n{a}\n{perm}"
    ia = [gen_data([2,2,2]), gen_data([2,3,4,5]), gen_data([5,3])]
    for a in ia: check(a)