jax.lax.platform_dependent doesn't stop Pallas from trying to lower for other backends?

[dlwh]

Description

I was trying to make a single function that supports CPU and non-CPU, and @froystig suggested I use jax.lax.platform_dependent. However, it doesn't seem to actually do the job (with pallas at least?)

from functools import partial

import jax
import jax.numpy as jnp
from jax.experimental.pallas.ops.tpu.megablox import gmm

key = jax.random.PRNGKey(42)

x = jax.random.normal(key, (128, 128))
weights = jax.random.normal(key, (3, 2, 128, 128))
group_sizes = jnp.array([60, 68])

def f(x, weight):
    out = jax.lax.platform_dependent(
        x, weight, group_sizes,
        tpu=partial(gmm, interpret=False),
        default=partial(gmm, interpret=True)
    )
    return out, out

out = jax.lax.scan(f, x, weights)

print(out)

...


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/dlwh/src/marin/tests/platform_dep.py", line 21, in <module>
    out = jax.lax.scan(f, x, weights)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/dlwh/src/marin/venv/lib/python3.11/site-packages/jax/_src/pallas/pallas_call.py", line 1300, in _pallas_call_lowering
    return mlir.lower_per_platform(ctx, "pallas_call",
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/dlwh/src/marin/venv/lib/python3.11/site-packages/jax/_src/pallas/pallas_call.py", line 1264, in cpu_lowering
    raise ValueError("Only interpret mode is supported on CPU backend.")
ValueError: Only interpret mode is supported on CPU backend.
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

System info (python version, jaxlib version, accelerator, etc.)

>>> import jax; jax.print_environment_info()
jax:    0.6.0
jaxlib: 0.6.0
numpy:  1.26.4
python: 3.11.11 (main, Dec  3 2024, 17:20:40) [Clang 16.0.0 (clang-1600.0.26.4)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='MacBook-Pro-57.local', release='24.4.0', version='Darwin Kernel Version 24.4.0: Fri Apr 11 18:33:47 PDT 2025; root:xnu-11417.101.15~117/RELEASE_ARM64_T6000', machine='arm64')

[blahBlahhhJ]

FYI without jax.lax.scan the code works fine.

[froystig]

@justinjfu and I partially debugged this...

The pallas_call lowering errs with the message about interpret mode whenever we lower a pallas_call on CPU without interpret mode.

In some contexts, we lower every one of the per-platform rules handed to lax.platform_dependent. Specifically, we emit an XLA switch and leave the right platform rule to be "selected as runtime" (more below).

We can illustrate this without lax.platform_dependent too, just a lax.cond:

from functools import partial
import jax, jax.numpy as jnp
from jax.experimental import pallas as pl

branch = partial(pl.pallas_call,
                 lambda _: None,
                 out_shape=jax.ShapeDtypeStruct((), jnp.int32))

def f(x):
  return x + jax.lax.cond(False,
                          lambda: branch(interpret=False)(),
                          lambda: branch(interpret=True)())

jax.jit(f)(1.)

We wrote "selected at runtime" in quotes above because, at some level, we actually know in principle what platform we're lowering for here. (That's because we're not only compiling for CPU, but also running on CPU, since we're using jax.jit, not an AOT or export path). Like the toy cond example above, you could say this is an issue where lowering (and the error-checking done in pallas_call lowering specifically) is too conservative or early.

[froystig]

FYI without jax.lax.scan the code works fine.

The immediate reason that scan contributes this in the example is that it causes us to stage out f via tracing. We can swap scan for jit in the example for a similar outcome. To see what I mean, look at the implementation of platform_dependent:

jax/jax/_src/lax/control_flow/conditionals.py

Lines 1076 to 1084 in 8137c37

	try:
	# Note/TODO(mvoz): This actually rarely seems to concretize - we could look into
	# core.ensure_compile_time_eval to get better single-branch selection.
	platform_index_concrete = core.concrete_or_error(operator.index, platform_index)
	except core.ConcretizationTypeError:
	return switch(platform_index, branches, *args)
	else:
	assert 0 <= platform_index_concrete < len(branches)
	return branches[platform_index_concrete](*args)

Under jit (or scan), platform_index_concrete isn't concrete, which leads us to produce the per-platform switch (i.e. take the except case) rather than inlining the CPU-specific platform branch (the else) case.

[gnecula]

I think this is a problem with how platform_dependent is implemented. It relies on the fact that all branches can actually be lowered on all platforms. Once there is a lowering, it relies on constant and conditional folding that happens early in compilation to keep only the branch we need (compilation is always for a single platform, while tracing and lowering can defer the choice of the platform).

In JAX we do have primitives that can only be lowered for some platforms, so platform_dependent should have a way to deal with such primitives.

I need to think some more about a good solution. I have tried a while ago to make platform_dependent a primitive, which means that it has its own lowering rules, and can be smart about only lowering the branches that it needs. But this is quite complicated because it will need to handle all the things that conditional primitives in JAX handle. Turning it into a switch seemed to be able to leverage all the features of switch. Except that now switch reasonably expects all branches to be lowerable.

[gnecula]

I have started to look into this and I created #28607 so far containing a reproducing test. As the test shows, lax.platform_dependent operates in a confusing way because it lowers all branches for the current lowering platform. This is all fine if branches are lowerable for all platforms, but I think that is too strong of a requirement. (I am surprised that we did not run into this before, but most JAX primitives are lowerable on all platforms, pallas_call and a few linalg primitives such as eig are rare exceptions.)

Solution A: The cleanest solution may be creating a conditional primitive for platform_dependent. Then its lowering rule can do the right thing and only lower the branches that are needed. But JAX conditional primitives are very complicated to implement. A variant of this solution is to add some bits from platform_dependent to cond_p, and then customize the lowering of cond_p for the case when it is a conditional on platform index. wdyt?

Solution B: We can try to resolve platform_index during tracing, so that we don't even introduce a conditional. This is what was attempted in

jax/jax/_src/lax/control_flow/conditionals.py

Lines 1076 to 1084 in 8137c37

	try:
	# Note/TODO(mvoz): This actually rarely seems to concretize - we could look into
	# core.ensure_compile_time_eval to get better single-branch selection.
	platform_index_concrete = core.concrete_or_error(operator.index, platform_index)
	except core.ConcretizationTypeError:
	return switch(platform_index, branches, *args)
	else:
	assert 0 <= platform_index_concrete < len(branches)
	return branches[platform_index_concrete](*args)

but as we see this does not work in a staging context (jit, cond, scan, etc.). Furthermore, so far the tracing in JAX is platform independent and I think it is best to leave it like that.

Solution C: We could do some constant and conditional folding during or just before lowering, when we know the actual platforms we are lowering for. It turns out that the Jaxpr that platform_dependent generates is always of the following form:

   platform_idx = platform_index (platforms=...)
   clamped_platform_idx = clamp 0 platform_idx 1
   res = cond branches=(branch_0, branch_1) clamped_platform_idx

We could pattern-match on this and remove unnecessary branches. The implementation is not too bad, but it would cost another Jaxpr traveral. (To check if we already have such a pre-lowering traversal.) Note that we have to handle both
the single-platform lowering when platform_idx is resolved to a constant, but also the multi-platform lowering.

During single-platform lowering the platform_index will be resolved to a constant (either 0 or 1). We'd have to constant fold the clamp and then when we lower the cond we can special case it for when the index is a constant. E.g., we have a platform_dependent branch that can only be lowered for TPU, and we are in multi-platform lowering mode for CPU and CUDA. The constant-folding should remove all irrelevant branches.

This pattern matching could be brittle to changes in how we trace, but I think we can catch those in tests. Maybe, there could be other code transformations that could legitimately change the Jaxpr so that our constant folding is foiled. This would not be the case for batching and AD because platform_index has no dependencies and appears as a constant.