code: def to_spec(tree):
      return jax.tree_util.tree_map(
          lambda x: ArraySpec.from_array(x)
          if isinstance(x, jax.Array)
          else x,
          tree,
      )
  
imports: from collections.abc import Mapping, Sequence
  import dataclasses
  import functools
  import pprint
  from typing import Any, Callable, Optional, TypeVar, Union
  from haiku._src import base
  from haiku._src import data_structures
  from haiku._src import module
  from haiku._src import transform
  from haiku._src import utils
  import jax
  import jax.numpy
  import numpy
  import tabulate
label_desc: Use image node_id 2 for calling a global function with example usage: to_spec(tree) and returns: jax