Flatten/unflatten custom Pytree before primitive bind_p call?

[femtomc]

(Edit: self-contained MWE at the bottom)

I'm working on a custom language which includes a set of primitives which I desugar to JAX understood constructs.

I started testing programs which accept Pytree instances as arguments, here is an example:

def hidden_markov_model(key, T, config):
    z0 = initial_position(config)
    key, z = trace("z", kernel)(key, (T, z0, config))
    return key, z

Here, config is a Pytree -- and not a simple type like a JAX array.

trace is a primitive defined in my system.

When I attempt to execute this through some of my API functions, I get the following error:

  File "/Users/mccoybecker/research/genjax/genjax/builtin/intrinsics.py", line 48, in <lambda>
    return lambda key, args: _trace(
  File "/Users/mccoybecker/research/genjax/genjax/builtin/intrinsics.py", line 38, in _trace
    return gen_fn_p.bind(
TypeError: Value DiscreteHMMConfiguration(linear_grid_dim=10, adjacency_distance_trans=1, adjacency_distance_obs=1, sigma_trans=0.1, sigma_obs=0.1, transition_tensor=Traced<ShapedArray(float32[10,10])>with<DynamicJaxprTrace(level=1/0)>, observation_tensor=Traced<ShapedArray(float32[10,10])>with<DynamicJaxprTrace(level=1/0)>) with type <class '__main__.DiscreteHMMConfiguration'> is not a valid JAX type

E.g. in the bind_p call to my trace primitive (trace is sugar to gen_fn_p), there's a call to concrete_aval for the config argument that throws. This I understand — but can I get around this?

I handle the primitives later (using a special transformation) which removes them from code — so really they are just sort of used as indicator sites for the later transformation.

More generally: if an object is a Pytree, is there something special I should be doing before/after bind_p calls? What's the right idiom for this?

Edit: here is a completely self-contained MWE demonstrating the thing I'm running into:

import jax


@jax.tree_util.register_pytree_node_class
class Foo:
    def tree_flatten(self):
        return (), ()

    @classmethod
    def tree_unflatten(cls, xs, data):
        return Foo()


g_p = jax.core.Primitive("some_g")


def g(*args):
    return g_p.bind(*args)


def bar(foo):
    return g(foo)


jaxpr = jax.make_jaxpr(bar)(Foo())
print(jaxpr)

Traceback (most recent call last):
  File "/Users/mccoybecker/research/genjax/examples/mwe_primitive_fail.py", line 25, in <module>
    jaxpr = jax.make_jaxpr(bar)(Foo())
  File "/Users/mccoybecker/research/genjax/examples/mwe_primitive_fail.py", line 22, in bar
    return g(foo)
  File "/Users/mccoybecker/research/genjax/examples/mwe_primitive_fail.py", line 18, in g
    return g_p.bind(*args)
TypeError: Value <__main__.Foo object at 0x135599d80> with type <class '__main__.Foo'> is not a valid JAX type

Should my "primitive wrapper" function g flatten/unflatten the Pytree at the bind boundary?

[femtomc]

To support this, I did exactly what I indicated with my question -- flatten/unflatten across bind boundaries.

A simple idiom: attach the tree_form from jax.tree_util.tree_flatten as metadata to the primitive (as a keyword argument to bind).