code: def get_ref_state_effects(
      avals: Sequence[core.AbstractValue], effects: core.Effects
  ) -> list[set[StateEffect]]:
      return [
          {
              eff
              for eff in effects
              if isinstance(eff, (ReadEffect, WriteEffect, AccumEffect))
              and eff.input_index == i
          }
          for i, _ in enumerate(avals)
      ]
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import dataclasses
  import math
  from typing import Any, Generic, TypeVar, Union
  from jax._src import core
  from jax._src import effects
  from jax._src import pretty_printer
  from jax._src.state import indexing
  from jax._src.util import safe_map, safe_zip
  from jax._src.typing import Array
label_desc: Use image node_id 4 for calling a global function with example usage: get_ref_state_effects(avals, effects) and returns: unknown