code: def _split_on_trainable(opt_state):
      trainable = {}
      non_trainable = {}
      for param_ref, value in opt_state.items():
          if param_ref.deref().trainable:
              trainable[param_ref] = value
          else:
              non_trainable[param_ref] = value
      return trainable, non_trainable
  
file_path: sonnet/sonnet/src/functional/optimizers.py
imports: import collections
  import functools
  from typing import Callable, Type
  from sonnet.src import base
  from sonnet.src.functional import haiku
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling a global function with example usage: _split_on_trainable(opt_state) and returns: trainable,  non_trainable