text: /**
   * @license
   * Copyright 2020 Google LLC. All Rights Reserved.
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   * http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   * =============================================================================
   */
  import { ENGINE } from '../engine';
  import { SpaceToBatchND } from '../kernel_names';
  import { convertToTensor } from '../tensor_util_env';
  import * as util from '../util';
  import { op } from './operation';
  /**
   * This operation divides "spatial" dimensions `[1, ..., M]` of the input into
   * a grid of blocks of shape `blockShape`, and interleaves these blocks with
   * the "batch" dimension (0) such that in the output, the spatial
   * dimensions `[1, ..., M]` correspond to the position within the grid,
   * and the batch dimension combines both the position within a spatial block
   * and the original batch position. Prior to division into blocks,
   * the spatial dimensions of the input are optionally zero padded
   * according to `paddings`. See below for a precise description.
   *
   * ```js
   * const x = tf.tensor4d([1, 2, 3, 4], [1, 2, 2, 1]);
   * const blockShape = [2, 2];
   * const paddings = [[0, 0], [0, 0]];
   *
   * x.spaceToBatchND(blockShape, paddings).print();
   * ```
   *
   * @param x A `tf.Tensor`. N-D with `x.shape` = `[batch] + spatialShape +
   * remainingShape`, where spatialShape has `M` dimensions.
   * @param blockShape A 1-D array. Must have shape `[M]`, all values must
   * be >= 1.
   * @param paddings A 2-D array. Must have shape `[M, 2]`, all values must be >=
   *     0. `paddings[i] = [padStart, padEnd]` specifies the amount to zero-pad
   * from input dimension `i + 1`, which corresponds to spatial dimension `i`. It
   * is required that
   * `(inputShape[i + 1] + padStart + padEnd) % blockShape[i] === 0`
   *
   * This operation is equivalent to the following steps:
   *
   * 1. Zero-pad the start and end of dimensions `[1, ..., M]` of the input
   * according to `paddings` to produce `padded` of shape paddedShape.
   *
   * 2. Reshape `padded` to `reshapedPadded` of shape:
   * `[batch] + [paddedShape[1] / blockShape[0], blockShape[0], ...,
   * paddedShape[M] / blockShape[M-1], blockShape[M-1]] + remainingShape`
   *
   * 3. Permute dimensions of `reshapedPadded` to produce `permutedReshapedPadded`
   * of shape: `blockShape + [batch] + [paddedShape[1] / blockShape[0], ...,
   * paddedShape[M] / blockShape[M-1]] + remainingShape`
   *
   * 4. Reshape `permutedReshapedPadded` to flatten `blockShape` into the
   * batch dimension, producing an output tensor of shape:
   * `[batch * prod(blockShape)] + [paddedShape[1] / blockShape[0], ...,
   * paddedShape[M] / blockShape[M-1]] + remainingShape`
   *
   * @doc {heading: 'Tensors', subheading: 'Transformations'}
   */
  function spaceToBatchND_(x, blockShape, paddings) {
      const $x = convertToTensor(x, 'x', 'spaceToBatchND');
      util.assert($x.rank >= 1 + blockShape.length, () => `input rank ${$x.rank} should be > than [blockShape] ${blockShape.length}`);
      util.assert(paddings.length === blockShape.length, () => `paddings.shape[0] ${paddings.length} must be equal to [blockShape] ${blockShape.length}`);
      util.assert($x.shape.reduce((a, b, i) => {
          if (i > 0 && i <= blockShape.length) {
              return a &&
                  ((b + paddings[i - 1][0] + paddings[i - 1][1]) %
                      blockShape[i - 1] ===
                      0);
          }
          return a;
      }, true), () => `input spatial dimensions ${$x.shape.slice(1)} with paddings ${paddings.toString()} must be divisible by blockShapes ${blockShape.toString()}`);
      const forward = backend => backend.spaceToBatchND($x, blockShape, paddings);
      const inputs = { x: $x };
      const attrs = { blockShape, paddings };
      return ENGINE.runKernelFunc(forward, inputs, null /* gradient */, SpaceToBatchND, attrs);
  }
  export const spaceToBatchND = op({ spaceToBatchND_ });
  //# sourceMappingURL=space_to_batch_nd.js.map