code: def _project_on_columns(A, v):
      """
      Returns A.T.conj() @ v.
      """
      v_proj = tree_map(
          lambda X, y: _einsum("...n,...->n", X.conj(), y),
          A,
          v,
      )
      return tree_reduce(operator.add, v_proj)
  
imports: from functools import partial
  import operator
  import numpy
  import jax
  import jax.numpy
  from jax import device_put
  from jax import lax
  from jax import scipy
  from jax.tree_util import tree_leaves, tree_map, tree_structure, tree_reduce, Partial
  from jax._src import dtypes
  from jax._src.lax import lax
  from jax._src.util import safe_map
label_desc: Use image node_id 15 for calling a global function with example usage: _project_on_columns(A, v) and returns: tree_reduce