return tf.tidy(() => {

    const batchSize = boxPredictions.shape[0]

    let boxes = decodeBoxesLayer(
      tf.reshape(tf.tile(params.extra_dim, [batchSize, 1, 1]), [-1, 4]) as tf.Tensor2D,
      tf.reshape(boxPredictions, [-1, 4]) as tf.Tensor2D
    )
    boxes = tf.reshape(
      boxes,
      [batchSize, (boxes.shape[0] / batchSize), 4]
    )

    const scoresAndClasses = tf.sigmoid(tf.slice(classPredictions, [0, 0, 1], [-1, -1, -1]))
    let scores = tf.slice(scoresAndClasses, [0, 0, 0], [-1, -1, 1]) as tf.Tensor

    scores = tf.reshape(
      scores,
      [batchSize, scores.shape[1]]
    )

    const boxesByBatch = tf.unstack(boxes) as tf.Tensor2D[]
    const scoresByBatch = tf.unstack(scores) as tf.Tensor1D[]

    return {
      boxes: boxesByBatch,
      scores: scoresByBatch
    }

  })

    it('is padded to square by 2 columns', () => tf.tidy(() => {
      const imgTensor = tf.tensor4d(Array(24).fill(1), [1, 4, 2, 3])
      const result = padToSquare(imgTensor)

      expect(result.shape).toEqual([1, 4, 4, 3])

      const paddedCols = tf.unstack(result, 2)
      expect(paddedCols.length).toEqual(4)
      expect(paddedCols[0].dataSync()).toEqual(ones(12))
      expect(paddedCols[1].dataSync()).toEqual(ones(12))
      expect(paddedCols[2].dataSync()).toEqual(zeros(12))
      expect(paddedCols[3].dataSync()).toEqual(zeros(12))
    }))

function getCenterCoordinatesAndSizesLayer(x: tf.Tensor2D) {
  const vec = tf.unstack(tf.transpose(x, [1, 0]))

  const sizes = [
    tf.sub(vec[2], vec[0]),
    tf.sub(vec[3], vec[1])
  ]

  const centers = [
    tf.add(vec[0], tf.div(sizes[0], tf.scalar(2))),
    tf.add(vec[1], tf.div(sizes[1], tf.scalar(2)))
  ]

  return {
    sizes,
    centers
  }
}

 return tfc.tidy(() => {
   const axis = getParamValue('axis', node, tensorMap, context) as number;
   const tensor =
       getParamValue('tensor', node, tensorMap, context) as tfc.Tensor;
   return tfc.unstack(tensor, axis);
 });