def METHOD_NAME(shape, dim, descending = False):
    x = jt.random(shape)
    y, y_key = jt.argsort(x, dim=dim, descending=descending)
    loss = (y_key * y_key).sum()
    gs = jt.grad(loss, x)
    assert np.allclose(x.data*2, gs.data)