/**
 * @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 { MaxPoolWithArgmax } from '@tensorflow/tfjs-core';
import { backend_util } from '@tensorflow/tfjs-core';
import { assertNotComplex } from '../cpu_util';
import { maxPoolWithArgmaxImpl } from './MaxPoolWithArgmax_impl';
export const maxPoolWithArgmaxConfig = {
    kernelName: MaxPoolWithArgmax,
    backendName: 'cpu',
    kernelFunc: ({ inputs, attrs, backend }) => {
        const { x } = inputs;
        const { filterSize, strides, pad, includeBatchInIndex } = attrs;
        const cpuBackend = backend;
        assertNotComplex(x, 'MaxPoolWithArgmax');
        const values = cpuBackend.data.get(x.dataId).values;
        const convInfo = backend_util.computePool2DInfo(x.shape, filterSize, strides, [1, 1], pad);
        const [pooled, indexes] = maxPoolWithArgmaxImpl(values, x.shape, x.dtype, includeBatchInIndex, convInfo);
        const pooledDataId = cpuBackend.write(pooled, convInfo.outShape, x.dtype);
        const indexesDataId = cpuBackend.write(indexes, convInfo.outShape, x.dtype);
        return [
            { dataId: pooledDataId, shape: convInfo.outShape, dtype: x.dtype },
            { dataId: indexesDataId, shape: convInfo.outShape, dtype: 'int32' }
        ];
    }
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiTWF4UG9vbFdpdGhBcmdtYXguanMiLCJzb3VyY2VSb290IjoiIiwic291cmNlcyI6WyIuLi8uLi8uLi8uLi8uLi8uLi90ZmpzLWJhY2tlbmQtY3B1L3NyYy9rZXJuZWxzL01heFBvb2xXaXRoQXJnbWF4LnRzIl0sIm5hbWVzIjpbXSwibWFwcGluZ3MiOiJBQUFBOzs7Ozs7Ozs7Ozs7Ozs7R0FlRztBQUVILE9BQU8sRUFBQyxpQkFBaUIsRUFBa0QsTUFBTSx1QkFBdUIsQ0FBQztBQUN6RyxPQUFPLEVBQUMsWUFBWSxFQUEyQixNQUFNLHVCQUF1QixDQUFDO0FBRTdFLE9BQU8sRUFBQyxnQkFBZ0IsRUFBQyxNQUFNLGFBQWEsQ0FBQztBQUU3QyxPQUFPLEVBQUMscUJBQXFCLEVBQUMsTUFBTSwwQkFBMEIsQ0FBQztBQUUvRCxNQUFNLENBQUMsTUFBTSx1QkFBdUIsR0FBaUI7SUFDbkQsVUFBVSxFQUFFLGlCQUFpQjtJQUM3QixXQUFXLEVBQUUsS0FBSztJQUNsQixVQUFVLEVBQUUsQ0FBQyxFQUFDLE1BQU0sRUFBRSxLQUFLLEVBQUUsT0FBTyxFQUFDLEVBQUUsRUFBRTtRQUN2QyxNQUFNLEVBQUMsQ0FBQyxFQUFDLEdBQUcsTUFBaUMsQ0FBQztRQUM5QyxNQUFNLEVBQUMsVUFBVSxFQUFFLE9BQU8sRUFBRSxHQUFHLEVBQUUsbUJBQW1CLEVBQUMsR0FDakQsS0FBMEMsQ0FBQztRQUMvQyxNQUFNLFVBQVUsR0FBRyxPQUF5QixDQUFDO1FBQzdDLGdCQUFnQixDQUFDLENBQUMsRUFBRSxtQkFBbUIsQ0FBQyxDQUFDO1FBRXpDLE1BQU0sTUFBTSxHQUFHLFVBQVUsQ0FBQyxJQUFJLENBQUMsR0FBRyxDQUFDLENBQUMsQ0FBQyxNQUFNLENBQUMsQ0FBQyxNQUFvQixDQUFDO1FBQ2xFLE1BQU0sUUFBUSxHQUFHLFlBQVksQ0FBQyxpQkFBaUIsQ0FDM0MsQ0FBQyxDQUFDLEtBQXlDLEVBQUUsVUFBVSxFQUFFLE9BQU8sRUFDaEUsQ0FBQyxDQUFDLEVBQUUsQ0FBQyxDQUFDLEVBQUUsR0FBRyxDQUFDLENBQUM7UUFDakIsTUFBTSxDQUFDLE1BQU0sRUFBRSxPQUFPLENBQUMsR0FBRyxxQkFBcUIsQ0FDM0MsTUFBTSxFQUFFLENBQUMsQ0FBQyxLQUFLLEVBQUUsQ0FBQyxDQUFDLEtBQUssRUFBRSxtQkFBbUIsRUFBRSxRQUFRLENBQUMsQ0FBQztRQUU3RCxNQUFNLFlBQVksR0FDZCxVQUFVLENBQUMsS0FBSyxDQUFDLE1BQXNCLEVBQUUsUUFBUSxDQUFDLFFBQVEsRUFBRSxDQUFDLENBQUMsS0FBSyxDQUFDLENBQUM7UUFDekUsTUFBTSxhQUFhLEdBQ2YsVUFBVSxDQUFDLEtBQUssQ0FBQyxPQUFxQixFQUFFLFFBQVEsQ0FBQyxRQUFRLEVBQUUsQ0FBQyxDQUFDLEtBQUssQ0FBQyxDQUFDO1FBQ3hFLE9BQU87WUFDTCxFQUFDLE1BQU0sRUFBRSxZQUFZLEVBQUUsS0FBSyxFQUFFLFFBQVEsQ0FBQyxRQUFRLEVBQUUsS0FBSyxFQUFFLENBQUMsQ0FBQyxLQUFLLEVBQUM7WUFDaEUsRUFBQyxNQUFNLEVBQUUsYUFBYSxFQUFFLEtBQUssRUFBRSxRQUFRLENBQUMsUUFBUSxFQUFFLEtBQUssRUFBRSxPQUFPLEVBQUM7U0FDbEUsQ0FBQztJQUNKLENBQUM7Q0FDRixDQUFDIiwic291cmNlc0NvbnRlbnQiOlsiLyoqXG4gKiBAbGljZW5zZVxuICogQ29weXJpZ2h0IDIwMjAgR29vZ2xlIExMQy4gQWxsIFJpZ2h0cyBSZXNlcnZlZC5cbiAqIExpY2Vuc2VkIHVuZGVyIHRoZSBBcGFjaGUgTGljZW5zZSwgVmVyc2lvbiAyLjAgKHRoZSBcIkxpY2Vuc2VcIik7XG4gKiB5b3UgbWF5IG5vdCB1c2UgdGhpcyBmaWxlIGV4Y2VwdCBpbiBjb21wbGlhbmNlIHdpdGggdGhlIExpY2Vuc2UuXG4gKiBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXRcbiAqXG4gKiBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjBcbiAqXG4gKiBVbmxlc3MgcmVxdWlyZWQgYnkgYXBwbGljYWJsZSBsYXcgb3IgYWdyZWVkIHRvIGluIHdyaXRpbmcsIHNvZnR3YXJlXG4gKiBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiBcIkFTIElTXCIgQkFTSVMsXG4gKiBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC5cbiAqIFNlZSB0aGUgTGljZW5zZSBmb3IgdGhlIHNwZWNpZmljIGxhbmd1YWdlIGdvdmVybmluZyBwZXJtaXNzaW9ucyBhbmRcbiAqIGxpbWl0YXRpb25zIHVuZGVyIHRoZSBMaWNlbnNlLlxuICogPT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT1cbiAqL1xuXG5pbXBvcnQge01heFBvb2xXaXRoQXJnbWF4LCBNYXhQb29sV2l0aEFyZ21heEF0dHJzLCBNYXhQb29sV2l0aEFyZ21heElucHV0c30gZnJvbSAnQHRlbnNvcmZsb3cvdGZqcy1jb3JlJztcbmltcG9ydCB7YmFja2VuZF91dGlsLCBLZXJuZWxDb25maWcsIFR5cGVkQXJyYXl9IGZyb20gJ0B0ZW5zb3JmbG93L3RmanMtY29yZSc7XG5pbXBvcnQge01hdGhCYWNrZW5kQ1BVfSBmcm9tICcuLi9iYWNrZW5kX2NwdSc7XG5pbXBvcnQge2Fzc2VydE5vdENvbXBsZXh9IGZyb20gJy4uL2NwdV91dGlsJztcblxuaW1wb3J0IHttYXhQb29sV2l0aEFyZ21heEltcGx9IGZyb20gJy4vTWF4UG9vbFdpdGhBcmdtYXhfaW1wbCc7XG5cbmV4cG9ydCBjb25zdCBtYXhQb29sV2l0aEFyZ21heENvbmZpZzogS2VybmVsQ29uZmlnID0ge1xuICBrZXJuZWxOYW1lOiBNYXhQb29sV2l0aEFyZ21heCxcbiAgYmFja2VuZE5hbWU6ICdjcHUnLFxuICBrZXJuZWxGdW5jOiAoe2lucHV0cywgYXR0cnMsIGJhY2tlbmR9KSA9PiB7XG4gICAgY29uc3Qge3h9ID0gaW5wdXRzIGFzIE1heFBvb2xXaXRoQXJnbWF4SW5wdXRzO1xuICAgIGNvbnN0IHtmaWx0ZXJTaXplLCBzdHJpZGVzLCBwYWQsIGluY2x1ZGVCYXRjaEluSW5kZXh9ID1cbiAgICAgICAgYXR0cnMgYXMgdW5rbm93biBhcyBNYXhQb29sV2l0aEFyZ21heEF0dHJzO1xuICAgIGNvbnN0IGNwdUJhY2tlbmQgPSBiYWNrZW5kIGFzIE1hdGhCYWNrZW5kQ1BVO1xuICAgIGFzc2VydE5vdENvbXBsZXgoeCwgJ01heFBvb2xXaXRoQXJnbWF4Jyk7XG5cbiAgICBjb25zdCB2YWx1ZXMgPSBjcHVCYWNrZW5kLmRhdGEuZ2V0KHguZGF0YUlkKS52YWx1ZXMgYXMgVHlwZWRBcnJheTtcbiAgICBjb25zdCBjb252SW5mbyA9IGJhY2tlbmRfdXRpbC5jb21wdXRlUG9vbDJESW5mbyhcbiAgICAgICAgeC5zaGFwZSBhcyBbbnVtYmVyLCBudW1iZXIsIG51bWJlciwgbnVtYmVyXSwgZmlsdGVyU2l6ZSwgc3RyaWRlcyxcbiAgICAgICAgWzEsIDFdLCBwYWQpO1xuICAgIGNvbnN0IFtwb29sZWQsIGluZGV4ZXNdID0gbWF4UG9vbFdpdGhBcmdtYXhJbXBsKFxuICAgICAgICB2YWx1ZXMsIHguc2hhcGUsIHguZHR5cGUsIGluY2x1ZGVCYXRjaEluSW5kZXgsIGNvbnZJbmZvKTtcblxuICAgIGNvbnN0IHBvb2xlZERhdGFJZCA9XG4gICAgICAgIGNwdUJhY2tlbmQud3JpdGUocG9vbGVkIGFzIEZsb2F0MzJBcnJheSwgY29udkluZm8ub3V0U2hhcGUsIHguZHR5cGUpO1xuICAgIGNvbnN0IGluZGV4ZXNEYXRhSWQgPVxuICAgICAgICBjcHVCYWNrZW5kLndyaXRlKGluZGV4ZXMgYXMgSW50MzJBcnJheSwgY29udkluZm8ub3V0U2hhcGUsIHguZHR5cGUpO1xuICAgIHJldHVybiBbXG4gICAgICB7ZGF0YUlkOiBwb29sZWREYXRhSWQsIHNoYXBlOiBjb252SW5mby5vdXRTaGFwZSwgZHR5cGU6IHguZHR5cGV9LFxuICAgICAge2RhdGFJZDogaW5kZXhlc0RhdGFJZCwgc2hhcGU6IGNvbnZJbmZvLm91dFNoYXBlLCBkdHlwZTogJ2ludDMyJ31cbiAgICBdO1xuICB9XG59O1xuIl19