Pallas emulation

[jbuckman]

The Pallas documentation describes an "Emulation Mode": https://jax.readthedocs.io/en/latest/pallas/design.html#emulation-mode

It doesn't come with an example or any other reference, though. Is this a feature I could use right now? Or is it something that is supported by the design, but not yet implemented?

[stephen-huan]

Apologies for the noise, but I am also curious about emulating Pallas on CPU. Did you manage to figure it out? Thanks!

[sharadmv]

Try passing interpret=True into your Pallas call (i.e. pl.pallas_call(kernel, ..., interpret=True). This will emulate your kernel using HLO and should run on CPU. The performance, however, will likely not be very good.

[stephen-huan]

That works, thanks! Now I get what the error message and the interpret attribute means...

  File "/nix/store/mcia3zi2043p0rhl1f8g1yl99vd3wyj8-python3-3.12.5-env/lib/python3.12/site-packages/jax/_src/pallas/pallas_call.py", line 524, in _pallas_call_lowering
    raise ValueError("Only interpret mode is supported on CPU backend.")
ValueError: Only interpret mode is supported on CPU backend.

(function) def pallas_call(
    f: (...) -> None,
    out_shape: Any,
    *,
    grid_spec: GridSpec | None = None,
    debug: bool = False,
    grid: Grid | None = None,
    in_specs: Sequence[BlockSpec | NoBlockSpec] | NoBlockSpec = no_block_spec,
    out_specs: BlockSpec | NoBlockSpec | Sequence[BlockSpec | NoBlockSpec] = no_block_spec,
    input_output_aliases: dict[int, int] = {},
    interpret: bool = False,
    name: str | None = None,
    compiler_params: dict[str, Any] | None = None
) -> JitWrapped

I somehow missed pallas was documented in jax.experimental.pallas with the rest of jax; I was only reading the quickstart tutorial. Thanks and apologies for the noise again!