Minor issue, but I spent a while trying to understand why this line was here before realizing it was unnecessary.

Quick spot check, using this code:

from jax import config, jit, lax, partial, make_jaxpr
config.disable_omnistaging()
f = partial(lax.full, ())
print(make_jaxpr(f)(0))
print(repr(f(0)))
print(repr(jit(f)(0)))

Result on master:

{ lambda  ; a.
  let b = device_put a
  in (b,) }
DeviceArray(0, dtype=int32)
DeviceArray(0, dtype=int32)

Result on this branch:

{ lambda  ; a.
  let 
  in (a,) }
DeviceArray(0, dtype=int32)
DeviceArray(0, dtype=int32)

The key observation is that device arrays are being returned, despite the lack of an explicit device_put on this branch.