code: def test_ω_add_mul(getkey):
      # ω(...) initialisation
      ω = eqxi.ω
      a = [0, 1]
      b = [1, 2]
      c = (ω(a) + ω(b)).ω
      assert c == [1, 3]
  
      # ...**ω initialisation
      for treedef in treedefs:
          a = b = c = random_pytree(getkey(), treedef)
  
          e1 = (a**ω * 2 + b**ω * c**ω - 3).ω
          e2 = jtu.tree_map(
              lambda ai, bi, ci: ai * 2 + bi * ci - 3, a, b, c
          )
          assert tree_allclose(e1, e2)
  
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 1 for calling a global function with example usage: test_ω_add_mul(getkey) without return types