code: def segment_dim(inputs, dim, shapes):
      """Returns tuple of Tensors output from segmenting input Tensor along dim.
  
      The returned tuple of Tensors produced by 'segmenting' the Tensor along a
      certain dimension can be transformed to specified shapes.
  
      Example:
          input_tensor = tf.placeholder([2, 14, 3])
          one, two = segment_dim(input_tensor, dim=1,
                                 shapes=[TensorShape([3, 3]), TensorShape([5])])
          # one is a [2, 3, 3, 3] Tensor and two is a [2, 5, 3] Tensor.
  
      Args:
        inputs: `Tensor` to segment.
        dim: dimension of the Tensor to operate on. Negative numbers count back from
          the end of the dimensions.
        shapes: list of TensorShapes of the output 'segments' to produce.
  
      Returns:
        Tuple with resulting Tensors.
  
      Raises:
        ValueError: if the dim used at initialization is invalid. The valid range is
        (-d, d], where d is the number of dimensions of the input tensor.
      """
      inputs_shape = inputs.shape
      ndims = inputs_shape.ndims
      dynamic_shape = tf.shape(inputs)
      shape_as_list = [
          dynamic_shape[i] if s is None else s
          for i, s in enumerate(inputs_shape.as_list())
      ]
  
      if dim >= ndims or dim < -ndims:
          message = "Invalid dims ({:d})".format(dim)
          raise ValueError(message)
  
      pre_shape = shape_as_list[:dim]
      if dim == -1:
          post_shape = []
      else:
          post_shape = shape_as_list[(dim + 1) :]
  
      slice_begin = [0] * ndims
      slice_size = [-1] * ndims
  
      segments = []
      for shape in shapes:
          num_elements = shape.num_elements()
          slice_size[dim] = num_elements
          flat_slice = tf.slice(inputs, slice_begin, slice_size)
  
          final_shape = pre_shape + shape.as_list() + post_shape
          segments.append(tf.reshape(flat_slice, final_shape))
          slice_begin[dim] += num_elements
  
      return tuple(segments)
  
file_path: sonnet/sonnet/src/nets/dnc/util.py
imports: import numpy
  import tensorflow
  import tree
label_desc: Use image node_id 1 for calling a global function with example usage: segment_dim(inputs, dim, shapes) and returns: tuple