Return Box.Get from function?

[lockwo]

How should box.get be returned from a function? A naive approach fails

import jax
import jax.numpy as jnp
from jax._src.hijax import Box

box = Box(jnp.array([1.0, 2.0, 3.0]))

result = jax.vmap(lambda _: box.get())(jnp.arange(2))

yields 'BatchTrace' object has no attribute 'cur_qdd'

but with

from jax._src.interpreters import batching
batching.BatchTrace.cur_qdd = lambda self, x: x.cur_qdd()

it seems to work? I also tested this in more complicated batch norm examples, but I wanted to check what the official recommendation was since this seemed a bit simplistic.

[cgarciae]

Hey @lockwo, currently Mutable HiTypes (like Box) don't work with vmap or scan.