[BUG] genjax.mh and genjax.Map don't play well together

[midfield]

Describe the bug
See code below

To Reproduce

@genjax.gen
def model():
    loc = genjax.normal(0., 1.) @ 'loc'
    xs = genjax.Map(genjax.normal, in_axes=(None, 0))(loc, jnp.arange(10)) @ 'xs'
    return xs

_, trace = genjax.simulate(model)(key, ())

@genjax.gen
def proposal(choices):
    loc = choices['loc']
    xs = genjax.Map(genjax.normal, in_axes=(None, 0))(loc, jnp.arange(10)) @ 'xs'
    return xs

genjax.mh(proposal).apply(key, trace, ())

Expected behavior
Doesn't produce error.

Execution environment (please complete the following information):

• OS: Linux
• Python version: 3.11
• GenJAX version: 0.0.1

[femtomc]

@midfield I think the correct example is this:

@genjax.gen
def model():
    loc = genjax.normal(0., 1.) @ 'loc'
    xs = genjax.Map(genjax.normal, in_axes=(None, 0))(loc, jnp.arange(10)) @ 'xs'
    return xs

_, trace = genjax.simulate(model)(key, ())

@genjax.gen
def proposal(choices):
    loc = choices['loc']
    xs = genjax.Map(genjax.normal, in_axes=(None, 0))(loc, jnp.arange(10)) @ 'xs'
    return xs

genjax.mh(proposal).apply(key, trace, ())

(Note change to Map invocation in the proposal).

This throws for me:

ValueError: Expected tuple, got [Array(-0.10823099, dtype=float32), Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 
dtype=int32)].

with an error from one of the axis size dimension functions.

[midfield]

Yes, that's right, sorry for the typo. I edited the original report. But this is still an error no?

[femtomc]

It is! That's a type issue -- a function expects a tuple somewhere, and we're providing a list -- which is coming from some of the incremental computing tracers, etc.

I'll investigate and fix.