code: def test_ω_inplace(getkey):
      ω = eqxi.ω
      for treedef in treedefs:
          a = random_pytree(getkey(), treedef)
          b1 = ω(a).at[()].set(3).ω
          b2 = jtu.tree_map(lambda ai: ai.at[()].set(3), a)
          assert tree_allclose(b1, b2)
  
          a2 = jtu.tree_map(lambda x: x + 1, a)
  
          b3 = ω(a).at[()].set(ω(a2)).ω
          b4 = jtu.tree_map(
              lambda ai, a2i: ai.at[()].set(a2i[()]), a, a2
          )
          assert tree_allclose(b3, b4)
  
file_path: equinox/tests/test_misc.py
imports: import equinox.internal
  import jax
  import jax.numpy
  import jax.tree_util
  import pytest
  from .helpers import random_pytree, tree_allclose, treedefs
label_desc: Use image node_id 2 for calling a global function with example usage: test_ω_inplace(getkey) without return types