code: def grid_env(env: tuple[tuple[Any, int], ...]) -> Iterator[None]:
      _grid_env_stack.append(
          tuple(
              GridEnv(axis_index, axis_size)
              for axis_index, axis_size in env
          )
      )
      try:
          yield
      finally:
          _grid_env_stack.pop()
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import contextlib
  import dataclasses
  import functools
  from typing import Any, Callable
  from collections.abc import Iterator
  from jax._src import api_util
  from jax._src import core
  from jax._src import linear_util
  from jax._src import state
  from jax._src import tree_util
  from jax._src import util
  from jax._src.interpreters import partial_eval
  from jax._src.state import discharge
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: grid_env(env) without return types