code: def shard_alike(x, y):
      """Shards x and y alike."""
      x_flat, x_tree = tree_flatten(x)
      y_flat, y_tree = tree_flatten(y)
  
      if x_tree != y_tree:
          raise ValueError(
              "Trees should be equal. "
              f"Got x_tree: {x_tree}, y_tree: {y_tree}"
          )
  
      for x_, y_ in safe_zip(x_flat, y_flat):
          x_aval = shaped_abstractify(x_)
          y_aval = shaped_abstractify(y_)
          if x_aval.shape != y_aval.shape:
              raise ValueError(
                  "The leaves shapes of `x` and `y` should match. Got `x` leaf shape:"
                  f" {x_aval.shape} and `y` leaf shape: {y_aval.shape}. File an issue at"
                  " https://github.com/google/jax/issues if you want this feature."
              )
  
      outs = [
          shard_alike_p.bind(x_, y_)
          for x_, y_ in safe_zip(x_flat, y_flat)
      ]
      x_out_flat, y_out_flat = zip(*outs)
      return tree_unflatten(x_tree, x_out_flat), tree_unflatten(
          y_tree, y_out_flat
      )
  
file_path: jax/jax/_src/shard_alike.py
imports: from functools import partial
  import itertools
  from jax._src import core
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.dispatch import apply_primitive
  from jax._src.tree_util import tree_flatten, tree_unflatten
  from jax._src.interpreters import batching
  from jax._src.util import safe_zip
  from jax._src.lib import xla_client
  from jax._src.api_util import shaped_abstractify
  from jax._src.lib.mlir import ir
label_desc: Use image node_id 1 for calling a global function with example usage: shard_alike(x, y) and returns: tree_unflatten,  tree_unflatten