/**
 * @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 { Pack, util } from '@tensorflow/tfjs-core';
import { concat } from './Concat';
import { expandDims } from './ExpandDims';
export function pack(args) {
    const { inputs, backend, attrs } = args;
    const { axis } = attrs;
    if (inputs.length === 1) {
        return expandDims({ inputs: { input: inputs[0] }, backend, attrs: { dim: axis } });
    }
    const shape = inputs[0].shape;
    const dtype = inputs[0].dtype;
    inputs.forEach(t => {
        util.assertShapesMatch(shape, t.shape, 'All tensors passed to stack must have matching shapes');
        util.assert(dtype === t.dtype, () => 'All tensors passed to stack must have matching dtypes');
    });
    const intermediateTensorInfos = [];
    const expandedTensors = inputs.map(t => {
        const expandedT = expandDims({ inputs: { input: t }, backend, attrs: { dim: axis } });
        intermediateTensorInfos.push(expandedT);
        return expandedT;
    });
    const result = concat({ inputs: expandedTensors, backend, attrs: { axis } });
    intermediateTensorInfos.forEach(t => backend.disposeIntermediateTensorInfo(t));
    return result;
}
export const packConfig = {
    kernelName: Pack,
    backendName: 'cpu',
    kernelFunc: pack
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiUGFjay5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uLy4uLy4uLy4uL3RmanMtYmFja2VuZC1jcHUvc3JjL2tlcm5lbHMvUGFjay50cyJdLCJuYW1lcyI6W10sIm1hcHBpbmdzIjoiQUFBQTs7Ozs7Ozs7Ozs7Ozs7O0dBZUc7QUFFSCxPQUFPLEVBQTJCLElBQUksRUFBcUMsSUFBSSxFQUFDLE1BQU0sdUJBQXVCLENBQUM7QUFHOUcsT0FBTyxFQUFDLE1BQU0sRUFBQyxNQUFNLFVBQVUsQ0FBQztBQUNoQyxPQUFPLEVBQUMsVUFBVSxFQUFDLE1BQU0sY0FBYyxDQUFDO0FBRXhDLE1BQU0sVUFBVSxJQUFJLENBQ2hCLElBQXFFO0lBRXZFLE1BQU0sRUFBQyxNQUFNLEVBQUUsT0FBTyxFQUFFLEtBQUssRUFBQyxHQUFHLElBQUksQ0FBQztJQUN0QyxNQUFNLEVBQUMsSUFBSSxFQUFDLEdBQUcsS0FBSyxDQUFDO0lBRXJCLElBQUksTUFBTSxDQUFDLE1BQU0sS0FBSyxDQUFDLEVBQUU7UUFDdkIsT0FBTyxVQUFVLENBQ2IsRUFBQyxNQUFNLEVBQUUsRUFBQyxLQUFLLEVBQUUsTUFBTSxDQUFDLENBQUMsQ0FBQyxFQUFDLEVBQUUsT0FBTyxFQUFFLEtBQUssRUFBRSxFQUFDLEdBQUcsRUFBRSxJQUFJLEVBQUMsRUFBQyxDQUFDLENBQUM7S0FDaEU7SUFFRCxNQUFNLEtBQUssR0FBRyxNQUFNLENBQUMsQ0FBQyxDQUFDLENBQUMsS0FBSyxDQUFDO0lBQzlCLE1BQU0sS0FBSyxHQUFHLE1BQU0sQ0FBQyxDQUFDLENBQUMsQ0FBQyxLQUFLLENBQUM7SUFFOUIsTUFBTSxDQUFDLE9BQU8sQ0FBQyxDQUFDLENBQUMsRUFBRTtRQUNqQixJQUFJLENBQUMsaUJBQWlCLENBQ2xCLEtBQUssRUFBRSxDQUFDLENBQUMsS0FBSyxFQUNkLHVEQUF1RCxDQUFDLENBQUM7UUFDN0QsSUFBSSxDQUFDLE1BQU0sQ0FDUCxLQUFLLEtBQUssQ0FBQyxDQUFDLEtBQUssRUFDakIsR0FBRyxFQUFFLENBQUMsdURBQXVELENBQUMsQ0FBQztJQUNyRSxDQUFDLENBQUMsQ0FBQztJQUVILE1BQU0sdUJBQXVCLEdBQWlCLEVBQUUsQ0FBQztJQUNqRCxNQUFNLGVBQWUsR0FBRyxNQUFNLENBQUMsR0FBRyxDQUFDLENBQUMsQ0FBQyxFQUFFO1FBQ3JDLE1BQU0sU0FBUyxHQUNYLFVBQVUsQ0FBQyxFQUFDLE1BQU0sRUFBRSxFQUFDLEtBQUssRUFBRSxDQUFDLEVBQUMsRUFBRSxPQUFPLEVBQUUsS0FBSyxFQUFFLEVBQUMsR0FBRyxFQUFFLElBQUksRUFBQyxFQUFDLENBQUMsQ0FBQztRQUNsRSx1QkFBdUIsQ0FBQyxJQUFJLENBQUMsU0FBUyxDQUFDLENBQUM7UUFDeEMsT0FBTyxTQUFTLENBQUM7SUFDbkIsQ0FBQyxDQUFDLENBQUM7SUFFSCxNQUFNLE1BQU0sR0FBRyxNQUFNLENBQUMsRUFBQyxNQUFNLEVBQUUsZUFBZSxFQUFFLE9BQU8sRUFBRSxLQUFLLEVBQUUsRUFBQyxJQUFJLEVBQUMsRUFBQyxDQUFDLENBQUM7SUFFekUsdUJBQXVCLENBQUMsT0FBTyxDQUMzQixDQUFDLENBQUMsRUFBRSxDQUFDLE9BQU8sQ0FBQyw2QkFBNkIsQ0FBQyxDQUFDLENBQUMsQ0FBQyxDQUFDO0lBRW5ELE9BQU8sTUFBTSxDQUFDO0FBQ2hCLENBQUM7QUFFRCxNQUFNLENBQUMsTUFBTSxVQUFVLEdBQWlCO0lBQ3RDLFVBQVUsRUFBRSxJQUFJO0lBQ2hCLFdBQVcsRUFBRSxLQUFLO0lBQ2xCLFVBQVUsRUFBRSxJQUE2QjtDQUMxQyxDQUFDIiwic291cmNlc0NvbnRlbnQiOlsiLyoqXG4gKiBAbGljZW5zZVxuICogQ29weXJpZ2h0IDIwMjAgR29vZ2xlIExMQy4gQWxsIFJpZ2h0cyBSZXNlcnZlZC5cbiAqIExpY2Vuc2VkIHVuZGVyIHRoZSBBcGFjaGUgTGljZW5zZSwgVmVyc2lvbiAyLjAgKHRoZSBcIkxpY2Vuc2VcIik7XG4gKiB5b3UgbWF5IG5vdCB1c2UgdGhpcyBmaWxlIGV4Y2VwdCBpbiBjb21wbGlhbmNlIHdpdGggdGhlIExpY2Vuc2UuXG4gKiBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXRcbiAqXG4gKiBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjBcbiAqXG4gKiBVbmxlc3MgcmVxdWlyZWQgYnkgYXBwbGljYWJsZSBsYXcgb3IgYWdyZWVkIHRvIGluIHdyaXRpbmcsIHNvZnR3YXJlXG4gKiBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiBcIkFTIElTXCIgQkFTSVMsXG4gKiBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC5cbiAqIFNlZSB0aGUgTGljZW5zZSBmb3IgdGhlIHNwZWNpZmljIGxhbmd1YWdlIGdvdmVybmluZyBwZXJtaXNzaW9ucyBhbmRcbiAqIGxpbWl0YXRpb25zIHVuZGVyIHRoZSBMaWNlbnNlLlxuICogPT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT1cbiAqL1xuXG5pbXBvcnQge0tlcm5lbENvbmZpZywgS2VybmVsRnVuYywgUGFjaywgUGFja0F0dHJzLCBQYWNrSW5wdXRzLCBUZW5zb3JJbmZvLCB1dGlsfSBmcm9tICdAdGVuc29yZmxvdy90ZmpzLWNvcmUnO1xuXG5pbXBvcnQge01hdGhCYWNrZW5kQ1BVfSBmcm9tICcuLi9iYWNrZW5kX2NwdSc7XG5pbXBvcnQge2NvbmNhdH0gZnJvbSAnLi9Db25jYXQnO1xuaW1wb3J0IHtleHBhbmREaW1zfSBmcm9tICcuL0V4cGFuZERpbXMnO1xuXG5leHBvcnQgZnVuY3Rpb24gcGFjayhcbiAgICBhcmdzOiB7aW5wdXRzOiBQYWNrSW5wdXRzLCBiYWNrZW5kOiBNYXRoQmFja2VuZENQVSwgYXR0cnM6IFBhY2tBdHRyc30pOlxuICAgIFRlbnNvckluZm8ge1xuICBjb25zdCB7aW5wdXRzLCBiYWNrZW5kLCBhdHRyc30gPSBhcmdzO1xuICBjb25zdCB7YXhpc30gPSBhdHRycztcblxuICBpZiAoaW5wdXRzLmxlbmd0aCA9PT0gMSkge1xuICAgIHJldHVybiBleHBhbmREaW1zKFxuICAgICAgICB7aW5wdXRzOiB7aW5wdXQ6IGlucHV0c1swXX0sIGJhY2tlbmQsIGF0dHJzOiB7ZGltOiBheGlzfX0pO1xuICB9XG5cbiAgY29uc3Qgc2hhcGUgPSBpbnB1dHNbMF0uc2hhcGU7XG4gIGNvbnN0IGR0eXBlID0gaW5wdXRzWzBdLmR0eXBlO1xuXG4gIGlucHV0cy5mb3JFYWNoKHQgPT4ge1xuICAgIHV0aWwuYXNzZXJ0U2hhcGVzTWF0Y2goXG4gICAgICAgIHNoYXBlLCB0LnNoYXBlLFxuICAgICAgICAnQWxsIHRlbnNvcnMgcGFzc2VkIHRvIHN0YWNrIG11c3QgaGF2ZSBtYXRjaGluZyBzaGFwZXMnKTtcbiAgICB1dGlsLmFzc2VydChcbiAgICAgICAgZHR5cGUgPT09IHQuZHR5cGUsXG4gICAgICAgICgpID0+ICdBbGwgdGVuc29ycyBwYXNzZWQgdG8gc3RhY2sgbXVzdCBoYXZlIG1hdGNoaW5nIGR0eXBlcycpO1xuICB9KTtcblxuICBjb25zdCBpbnRlcm1lZGlhdGVUZW5zb3JJbmZvczogVGVuc29ySW5mb1tdID0gW107XG4gIGNvbnN0IGV4cGFuZGVkVGVuc29ycyA9IGlucHV0cy5tYXAodCA9PiB7XG4gICAgY29uc3QgZXhwYW5kZWRUID1cbiAgICAgICAgZXhwYW5kRGltcyh7aW5wdXRzOiB7aW5wdXQ6IHR9LCBiYWNrZW5kLCBhdHRyczoge2RpbTogYXhpc319KTtcbiAgICBpbnRlcm1lZGlhdGVUZW5zb3JJbmZvcy5wdXNoKGV4cGFuZGVkVCk7XG4gICAgcmV0dXJuIGV4cGFuZGVkVDtcbiAgfSk7XG5cbiAgY29uc3QgcmVzdWx0ID0gY29uY2F0KHtpbnB1dHM6IGV4cGFuZGVkVGVuc29ycywgYmFja2VuZCwgYXR0cnM6IHtheGlzfX0pO1xuXG4gIGludGVybWVkaWF0ZVRlbnNvckluZm9zLmZvckVhY2goXG4gICAgICB0ID0+IGJhY2tlbmQuZGlzcG9zZUludGVybWVkaWF0ZVRlbnNvckluZm8odCkpO1xuXG4gIHJldHVybiByZXN1bHQ7XG59XG5cbmV4cG9ydCBjb25zdCBwYWNrQ29uZmlnOiBLZXJuZWxDb25maWcgPSB7XG4gIGtlcm5lbE5hbWU6IFBhY2ssXG4gIGJhY2tlbmROYW1lOiAnY3B1JyxcbiAga2VybmVsRnVuYzogcGFjayBhcyB1bmtub3duIGFzIEtlcm5lbEZ1bmNcbn07XG4iXX0=