Non-hashable static arguments not allowed for pmap but allowed for vmap?

[somearthling]

I have this piece of code

qnode = jit(qml.QNode(circuit, dev, interface='jax'), static_argnums=2)
probs = pmap(qnode, in_axes=(0, None, None))(x, p, filt)

inside a function that is already jitted, and so x and p are jax tracers at this point. This runs completely fine, but when I replace the vmap with a pmap, I get the following error:

ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses. Static argument (index 2) of type <class 'jax.interpreters.partial_eval.DynamicJaxprTracer'> for function circuit is non-hashable.

Is this intended, or is there something I'm doing wrong that I can fix?

[jakevdp]

This sounds like a bug: a traced value should never be valid as a static argument, so you should get an error regardless of whether you use pmap or vmap. That said, I'm having trouble reproducing the behavior you described. Can you put together a complete minimal example? Here's what I tried:

import os
os.environ['XLA_FLAGS'] = " --xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp

def f(x, y):
  return x + y

@jax.jit
def g_vmap(x, y):
  jax.vmap(jax.jit(f, static_argnums=1))(x, y)
  
@jax.jit
def g_pmap(x, y):
  jax.pmap(jax.jit(f, static_argnums=1))(x, y)

x = jnp.ones((8, 2))
y = jnp.ones((8, 2))

g_vmap(x, y)
# ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses.

g_pmap(x, y)
# ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses.

[somearthling]

Sorry, I think there was a typo in my code, the one that runs fine is vmap. It returns the error when I switch to pmap. In this case, x and p are non-hashable, and filt is not - I've marked it out as a static argument everywhere. Not sure why it says argument of index 2 is non-hashable. This should be filt, correct? Your code is a little different from mine, though, I have a function gradient and cost:

@partial(jit, static_argnums=(3,))
def gradient_and_cost(x, y, p, filt):
  ...
  qnode = jit(qml.QNode(circuit, dev, interface='jax'), static_argnums=2)
  probs = vmap(qnode, in_axes=(0, None, None))(x, p, filt)

I'm thinking this has something to do with the partial? qnode should output a single vector of length 4, in case that helps.

[jakevdp]

Thanks for the clarification – the issue here is that pmap jit-compiles its function by default, while vmap doesn't. If you want to use pmap in this fashion, you could close over the static arguments instead:

@partial(jit, static_argnums=(3,))
def gradient_and_cost(x, y, p, filt):
  ...
  # you may have to modify this partial argument depending on the signature of QNotde.__call__:
  qnode = jit(partial(qml.QNode(circuit, dev, interface='jax'), filt=filt))
  probs = pmap(qnode, in_axes=(0, None))(x, p)