Why does vmap not preserve static (concrete) arguments in jit?

[femtomc]

I have some code which passes a static argument into a jit call. Before the vmap call, everything is fine:

JAXChoiceMap(
  trie=Trie(
    nodes={
      'z':
      Mask(
        inner=VectorChoiceMap(
          subtrace=JAXChoiceMap(
            trie=Trie(nodes={'obs': ValueChoiceMap(value=f32[1,8,2](numpy))})
          )
        ),
        mask=VectorChoiceMap(
          subtrace=JAXChoiceMap(
            trie=Trie(nodes={'obs': ValueChoiceMap(value=bool[1,8](numpy))})
          )
        )
      )
    }
  )
)

Here, the concrete arrays I created are untouched inside the JIT call.

However, when I pass this into vmap and choose not to broadcast (by setting in_axes = (..., None) in the appropriate spot), I get the following:

JAXChoiceMap(
  trie=Trie(
    nodes={
      'z':
      Mask(
        inner=VectorChoiceMap(
          subtrace=JAXChoiceMap(
            trie=Trie(nodes={'obs': ValueChoiceMap(value=f32[1,8,2])})
          )
        ),
        mask=VectorChoiceMap(
          subtrace=JAXChoiceMap(
            trie=Trie(nodes={'obs': ValueChoiceMap(value=bool[1,8])})
          )
        )
      )
    }
  )
)

Now, this concrete Pytree has had its leaves converted to traced values.

Is this intended behavior? From prior discussion: #10712 it seems like vmap should leave the leaves alone on this Pytree instance -- or am I mistaken?

Note: to try and resolve this, I also jax.tree_util.tree_map over the Pytree, created a "None"-tree and passed that in as in_axes for the Pytree slot. This also did not prevent JAX from lifting the numpy arrays to traced values.

[jakevdp]

I'm having difficulty understanding exactly what code you're executing, what the result is, and how it differs from what you expect. Could you edit your question to add a minimal reproducible example?

[jakevdp]

You should be able to do this by passing in_axes=None for the value you want to be untouched. Here is a short example:

from typing import NamedTuple
import jax.numpy as jnp
import jax

# Note: NamedTuple is automatically registered as a PyTree
class Data(NamedTuple):
  x: jnp.ndarray
  y: int

data = Data(jnp.arange(5), 2)

def f(data):
  assert isinstance(data.y, int)  # asserts that `data.y` is static/unchanged
  return data.x + data.y

out = jax.vmap(f, in_axes=[Data(0, None)])(data)
print(out)
# [2 3 4 5 6]

If that's not working in your case, it's difficult to help without a reproducible example of the behavior you're seeing.

[femtomc]

Thank you -- the reason why I can't provide a ready example is that it's in the middle of a rather complicated set of computations. I may have to share a larger Gist, if that's okay?

[femtomc]

But also -- I think you've indicated that I've made a mistake in specifying where not to broadcast, which is probably correct :) So also thank you for that.

[femtomc]

Yes -- it turned out to be an intermediate transformation I was performing (lifted each concrete value to a traced one, my own ignorance).

[jakevdp]

Great! Glad the example helped 😁