code: def assert_same_shallow_tree(shallow, tree):
      """Asserts that `tree` has the same shallow structure as `shallow`."""
      # Do a dummy map for the side-effect of verifying that the structures are
      # the same. This doesn't catch all the errors we actually care about, sadly.
      map_tree_up_to(shallow, lambda *args: (), tree)
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  from jax import tree_util
  import jax.numpy
label_desc: Use image node_id 6 for calling a global function with example usage: assert_same_shallow_tree(shallow, tree) without return types