code: def shaped_allclose(x, y, **kwargs):
      if isinstance(x, jax.Array):
          x = jax.device_put(x)
      return _shaped_allclose(x, y, **kwargs)
  
imports: import functools
  from typing import Any, Union
  import equinox
  import jax
  import jax.numpy
  import jax.random
  import jax.tree_util
  import pytest
  from .helpers import tree_allclose
label_desc: Use image node_id 1 for calling a global function with example usage: shaped_allclose(x, y) and returns: _shaped_allclose