code: def _read_keys(key, x1, x2):
      """Read dropout key.
  
      `key` might be a tuple of two rng keys or a single rng key or None. In
      either case, `key` will be mapped into two rng keys `key1` and `key2` to
      make sure `(x1==x2) == (key1==key2)`.
      """
  
      if key is None or all_none(x2):
          key1 = key2 = key
      elif isinstance(key, tuple) and len(key) == 2:
          key1, key2 = key
          new_key = jnp.where(
              x1_is_x2(key1, key2), random.fold_in(key2, 1), key2
          )
          key2 = jnp.where(x1_is_x2(x1, x2), key1, new_key)
          warnings.warn(
              "The value of `key[1]` might be replaced by a new value if "
              "key[0] == key[1] and x1 != x2 or key[0] != key[1] and "
              "x1 == x2."
          )
      elif isinstance(key, jnp.ndarray):
          key1 = key
          key2 = jnp.where(
              x1_is_x2(x1, x2), key1, random.fold_in(key, 1)
          )
      else:
          raise TypeError(type(key))
      return key1, key2
  
imports: from collections import namedtuple
  import functools
  import inspect
  import operator
  import types
  from typing import Any, Callable, Iterable, Optional, Sequence, Sized, TypeVar, Union
  import warnings
  import jax
  from jax import core
  from jax import random
  import jax.numpy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_map
  import numpy
label_desc: Use image node_id 26 for calling a global function with example usage: _read_keys(key, x1, x2) and returns: key1,  key2