code: def PolicyAndValue(
      policy_distribution,
      body=None,
      policy_top=Policy,
      value_top=Value,
      normalizer=None,
      joint=True,
      mode="train",
  ):
      """Attaches policy and value heads to a model body."""
      if normalizer is None:
          normalizer = lambda mode: []
      if body is None:
          body = lambda mode: []
  
      common_kwargs = {"body": None, "normalizer": None, "mode": mode}
      policy_top = policy_top(
          policy_distribution=policy_distribution, **common_kwargs
      )
      value_top = value_top(**common_kwargs)
  
      layers = [normalizer(mode=mode)]
      if joint:
          layers.extend(
              [
                  body(mode=mode),
                  tl.Branch(policy_top, value_top),
              ]
          )
      else:
          layers.append(
              tl.Branch(
                  tl.Serial(body(mode=mode), policy_top),
                  tl.Serial(body(mode=mode), value_top),
              )
          )
      return tl.Serial(layers)
  
imports: from trax import layers
  from trax import models
  from trax.fastmath import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: PolicyAndValue(policy_distribution, body, policy_top, value_top, normalizer, joint, mode) and returns: tl