jax.jit(foo, device=dev) vs jax.device_put and interpretted vs JIT-ed

[erick-xanadu]

Hi,

I am thinking about tentatively implementing a PJRT plugin and I have some questions about the interfaces for users. At the moment, I imagine my PJRT plugin as taking as an input an MLIR module via this function:

  // Variant of `Compile` that accepts an MLIR module.
  virtual absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
      mlir::ModuleOp module, CompileOptions options) = 0;

compiling the module to some executable and then later receiving the concrete inputs and executing the function. I really like jax.jit(foo, device=dev) but I notice that there are some comments that the API might change. There's also jax.device_put.

I think the following two are equivalent but I am not sure. Can you confirm?

x = jnp.array(0)
jax.jit(foo, device=dev)(x)

and

x = jnp.array(0)
x = jax.device_put(x, dev)
jax.jit(foo)(x)

I can see this working by jax.jit just determining the target by looking at all inputs current device placement.

Also, I have a question about backends when things are not executed via jax.jit. If I execute:

# no jax.jit context
x = jnp.array(0)
x = jax.device_put(x, dev)
x + x

Does this mean that the python line x + x will generate an MLIR module containing a simple stablehlo.addOp and returning it? And then using the PJRT plugin as in the case of jax.jit? Or how would the execution pipeline look like? Thanks!

[yashk2810]

jax.jit(foo, device=dev)(x) vs x = jax.device_put(x, dev); jax.jit(foo)(x)

Yes, these should be equivalent but only in single device case. If x is sharded, then they won't be equivalent. Also jax.jit(f, device=) is deprecated so you should use the latter.

I can see this working by jax.jit just determining the target by looking at all inputs current device placement.

Yes, that's correct. JAX uses computation follows data semantics.

Does this mean that the python line x + x will generate an MLIR

+ is jitted internally, so it will be executed on the default device on the default backend that's present.

[erick-xanadu]

+ is jitted internally, so it will be executed on the default device on the default backend that's present.

Even if I specify the device for x it will be executed on the default device? I think the answer is essentially "yes the addition will be JIT compiled by whatever device x is placed on"?

[yashk2810]

Nope, then we use the computation follows data semantics to execute on the device x is on. See this: https://jax.readthedocs.io/en/latest/faq.html#controlling-data-and-computation-placement-on-devices

Yeah, we are on the same page :)