def dummy_root_usage(x):
    f = (lambda y: (x - y))
    return lax.custom_root(f, 0.0, (lambda f, x: x), (lambda f, x: (x, x)))
def dummy_root_usage(x):
    f = (lambda y: (x - y))
    return lax.custom_root(f, 0.0, (lambda f, x: x), (lambda f, x: (x, x)))