[platform_dependent] Ensure that platform_dependent only lowers for intended platforms

Fixes: #28594

Currently `lax.platform_dependent` allows specifying code that behaves
differently when lowered on different platforms. However, this function
operates in a confusing way, in that it will create a branch on the
platform, but will lower all branches for the **current** lowering platforms.

For example, in the following code:
```
   lax.platform_dependent(x,
                          cpu=for_cpu, tpu=for_tpu)
```

If we lower for CPU, we lower both `for_cpu` and `for_tpu`
for CPU (!), but only the branch corresponding to `for_cpu`
will actually run.

This is a problem if, e.g., `for_tpu` does not have a lowering
for CPU. We will get an error during lowering. Instead there should
be no error during lowering, because that branch is not actually needed.

We add a new test `test_platform_dependent_with_primitive_with_lowering_error`
to demonstrate this.

The solution implememented here is the Solution A from #28594: we
add a `branches_platform` param to the `cond` primitive, which is
propagated by all transformations. This param is used only for the
conditionals arising from `lax.platform_dependendet`.
During lowering we drop the branches corresponding to the platforms
that are not interesting.

Author: gnecula

jax/_src/checkify.py

@@ -759,7 +759,8 @@ def jaxpr_to_checkify_jaxpr(
   out_tree, error_effects = metadata()
   return checked_jaxpr, out_tree, error_effects
 
-def cond_error_check(error: Error, enabled_errors, index, *ops, branches):
+def cond_error_check(error: Error, enabled_errors, index, *ops,
+                     branches, **params):
   # Get the error-effects out of all branches so the cond can be called with
   # a merged error with all these effects.
   err_vals, err_tree = jtu.tree_flatten(error)
@@ -780,7 +781,7 @@ def get_error_effects_from_jaxpr(jxpr):
 
   err_and_outs = lax.cond_p.bind(
       index, *err_vals, *ops,
-      branches=tuple(new_branches))
+      branches=tuple(new_branches), **params)
 
   # we need to merge metadata across out_trees (a tuple)
   err0, out = tree_unflatten(out_trees[0], err_and_outs)

jax/_src/interpreters/mlir.py

@@ -2080,6 +2080,11 @@ def _platforms_for_eqn_ctx(eqn_ctx: core.JaxprEqnContext | None
     return ('tpu',)
   return ()
 
+def _platforms_for_eqn(ctx: LoweringRuleContext) -> tuple[str, ...]:
+  """The lowering platforms for the current eqn"""
+  return tuple((_platforms_for_eqn_ctx(ctx.jaxpr_eqn_ctx) or
+               ctx.platforms or ctx.module_context.platforms))
+
 
 def lower_per_platform(ctx: LoweringRuleContext,
                        description: str,
@@ -2122,8 +2127,7 @@ def lower_per_platform(ctx: LoweringRuleContext,
    rule_args: the args of the lowering rules.
    rule_kwargs: the kwargs of the lowering rules.
   """
-  platforms: Sequence[str] = (_platforms_for_eqn_ctx(ctx.jaxpr_eqn_ctx) or
-                              ctx.platforms or ctx.module_context.platforms)
+  platforms: Sequence[str] = _platforms_for_eqn(ctx)
   # Special case the common case (single-platform lowering)
   if len(platforms) == 1:
     rule = platform_rules.get(platforms[0], default_rule)