code: def policy_action(
      apply_fn: Callable[..., Any],
      params: flax.core.frozen_dict.FrozenDict,
      state: np.ndarray,
  ):
      """Forward pass of the network.
  
      Args:
        params: the parameters of the actor-critic model
        module: the actor-critic model
        state: the input for the forward pass
  
      Returns:
        out: a tuple (log_probabilities, values)
      """
      out = apply_fn({"params": params}, state)
      return out
  
file_path: flax/examples/ppo/agent.py
imports: import collections
  import functools
  import multiprocessing
  from typing import Any, Callable
  import flax
  import jax
  import numpy
  import env_utils
label_desc: Use image node_id 1 for calling a global function with example usage: policy_action(apply_fn, params, state) and returns: out