code: def nested_stack(objs, axis=0, np_module=np):
      """Stacks the numpy arrays inside any dicts/lists/tuples in `objs`.
  
      Args:
        objs: List of nested structures to stack.
        axis: Axis to stack along.
        np_module: numpy module to use - typically numpy or jax.numpy.
  
      Returns:
        An object with the same nested structure as each element of `objs`, with
        leaves stacked together into numpy arrays. Nones are propagated, i.e. if
        each element of the stacked sequence is None, the output will be None.
      """
      # nested_map the stacking operation, but stopping at level 1 so at tuples of
      # numpy arrays.
      return nested_map(
          lambda x: np_module.stack(x, axis=axis),
          nested_zip(objs),
          level=1,
      )
  
file_path: trax/trax/fastmath/numpy.py
imports: import numpy
  from scipy.special import logsumexp
  from trax.shapes import signature
label_desc: Use image node_id 10 for calling a global function with example usage: nested_stack(objs, axis, np_module) and returns: nested_map