def handling(
      cls,
      body: layer_base.LayerLike,
      initial_key: jax.Array,
      hole_predicate: Callable[
          [RandomRequest | TaggedRandomRequest], bool
      ] = _is_untagged_hole,
      state_category: Any = "random",
      handler_id: str | None = None,
  ) -> WithStatefulRandomKey:
    """Builds a `WithStatefulRandomKey` that handles effects in this layer.

    Args:
      body: The layer to wrap. Usually will contain random effects in the form
        of `RandomRequest` or `TaggedRandomRequest`.
      initial_key: Initial key to use for the state.
      hole_predicate: Callable that determines whether we should handle a given
        random effect hole. By default, handles all instances of `RandomRequest`
        but no instances of `TaggedRandomRequest`.
      state_category: Type to use when configuring the state effect.
      handler_id: ID to use for the handler. If None, will be inferred.

    Returns:
      A `WithStatefulRandomKey` handler wrapping `body`, with its random
      effect holes replaced with references to this handler (whenever allowed
      by the predicate).
    """
    shapecheck.check_structure(
        initial_key,
        shapecheck.ArraySpec(shape=(), dtype=jax.dtypes.prng_key),
    )
    handler_id = effect_base.infer_or_check_handler_id(
        "WithStatefulRandomKey", body, explicit_id=handler_id
    )
    ref = HandledRandomRef(handler_id=handler_id)
    adjusted_body = (
        selectors.select(body)
        .at_instances_of(RandomRequest | TaggedRandomRequest)
        .where(hole_predicate)
        .set(ref)
    )
    return cls(
        handler_id=handler_id,
        body=adjusted_body,
        random_state=local_state.FrozenLocalStateRequest(
            state=initial_key, category=state_category
        ),
    )
Builds a `WithStatefulRandomKey` that handles effects in this layer.

    Args:
      body: The layer to wrap. Usually will contain random effects in the form
        of `RandomRequest` or `TaggedRandomRequest`.
      initial_key: Initial key to use for the state.
      hole_predicate: Callable that determines whether we should handle a given
        random effect hole. By default, handles all instances of `RandomRequest`
        but no instances of `TaggedRandomRequest`.
      state_category: Type to use when configuring the state effect.
      handler_id: ID to use for the handler. If None, will be inferred.

    Returns:
      A `WithStatefulRandomKey` handler wrapping `body`, with its random
      effect holes replaced with references to this handler (whenever allowed
      by the predicate).
    
penzai/deprecated/v1/data_effects/random.py