code: def test_split_outputs_per_head():
      class TestHead:
          def __init__(self, num_outputs):
              self.num_outputs = num_outputs
  
      assert split_outputs_per_head(
          (1, 2, 3, 4), [TestHead(2), TestHead(1), TestHead(1)]
      ) == [[1, 2], [3], [4]]
  
imports: import pytest
  import mxnet
  from mxnet import nd
  import numpy
  from rl_coach.architectures.mxnet_components.utils import *
label_desc: Use image node_id 4 for calling a global function with example usage: test_split_outputs_per_head() without return types