Pallas support on apple silicon metal

[s-zeng]

Hi all, I've tried to run the basic JAX demo on apple silicon, but am getting the following error:

In [4]: @jax.jit
   ...: def add_vectors(x: jax.Array, y: jax.Array) -> jax.Array:
   ...:   return pl.pallas_call(
   ...:       add_vectors_kernel,
   ...:       out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype)
   ...:   )(x, y)
   ...: 
In [6]: add_vectors(jnp.arange(10), jnp.arange(10))
---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation         Traceback (most recent call last)
...
NotImplementedError                       Traceback (most recent call last)
Cell In[6], line 1
----> 1 add_vectors(jnp.arange(10), jnp.arange(10))

    [... skipping hidden 13 frame]

File ~/repos/graphcast/.venv/lib/python3.12/site-packages/jax/_src/pallas/pallas_call.py:1518, in _pallas_call_lowering(ctx, interpret, backend, *in_nodes, **params)
   1512     raise _unsupported_lowering_error("gpu")
   1514   return pallas_call_registration.pallas_call_lowering(
   1515       ctx, *in_nodes, **params
   1516   )
-> 1518 return mlir.lower_per_platform(ctx, "pallas_call",
   1519                                dict(cpu=cpu_lowering,
   1520                                     tpu=tpu_lowering,
   1521                                     cuda=gpu_lowering,
   1522                                     rocm=gpu_lowering),
   1523                                None,  # default_rule
   1524                                effects.no_effects,
   1525                                *in_nodes,
   1526                                interpret=interpret,
   1527                                **params)

File ~/repos/graphcast/.venv/lib/python3.12/site-packages/jax/_src/interpreters/mlir.py:1996, in lower_per_platform(ctx, description, platform_rules, default_rule, effects, *rule_args, **rule_kwargs)
   1994   rule = platform_rules.get(platforms[0], default_rule)
   1995   if rule is None:
-> 1996     raise NotImplementedError(
   1997       f"MLIR translation rule for primitive '{description}' not "
   1998       f"found for platform {platforms[0]}")
   2000 # Multi-platform lowering
   2001 kept_rules: list[LoweringRule] = []  # Only the rules for the platforms of interest

MLIR translation rule for primitive 'pallas_call' not found for platform METAL

Is there a timeline for pallas support on apple silicon?

[dfm]

I'm not aware of efforts to support Pallas on metal, but @shuhand0 or @kulinseth might be able to comment more.