class_docstr: Provides friendlier syntax for mapping with `jax.tree_util.tree_map`.
  
  !!! example
  
      ```python
      (ω(a) + ω(b)).ω == jax.tree_util.tree_map(operator.add, a, b)
      ```
  
  !!! tip
  
      To minimise the number of brackets, the following `__rpow__` syntax can be
      used:
  
      ```python
      (a**ω + b**ω).ω == jax.tree_util.tree_map(operator.add, a, b)
      ```
  
      This is entirely equivalent to the above.
docstr: **Arguments:**
  
  - `value`: The PyTree to wrap.
  - `is_leaf`: An optional value for the `is_leaf` argument to
    `jax.tree_util.tree_map`.
  
  !!! note
  
      The `is_leaf` argument cannot be set when using the `__rpow__` syntax for
      initialisation.
docstr_tok: ["*","*","Arguments",":","*","*","-","`","value","`",":","The","PyTree","to","wrap",".","-","`","is_leaf","`",":","An","optional","value","for","the","`","is_leaf","`","argument","to","`","jax.tree_util.tree_map","`",".","!","!","!","note","The","`","is_leaf","`","argument","can","not","be","set","when","using","the","`","__rpow__","`","syntax","for","initialisation","."]
code: def __init__(self, value, is_leaf=None):
      """
      **Arguments:**
  
      - `value`: The PyTree to wrap.
      - `is_leaf`: An optional value for the `is_leaf` argument to
        `jax.tree_util.tree_map`.
  
      !!! note
  
          The `is_leaf` argument cannot be set when using the `__rpow__` syntax for
          initialisation.
      """
      self.ω = value
      self.is_leaf = is_leaf
  
imports: import operator
  from collections.abc import Callable
  from typing import Any, Optional, TYPE_CHECKING
  import jax
  import jax.tree_util
label_desc: Use image node_id 1 to create a new _ω object  with example: obj = _ω(value, is_leaf)