/**
 * @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 { BatchMatMul } from '../kernel_names';
import { matMul } from '../ops/mat_mul';
export const batchMatMulGradConfig = {
    kernelName: BatchMatMul,
    inputsToSave: ['a', 'b'],
    gradFunc: (dy, saved, attrs) => {
        const [a, b] = saved;
        const { transposeA, transposeB } = attrs;
        if (!transposeA && !transposeB) {
            return {
                a: () => matMul(dy, b, false, true),
                b: () => matMul(a, dy, true, false)
            };
        }
        else if (!transposeA && transposeB) {
            return {
                a: () => matMul(dy, b, false, false),
                b: () => matMul(dy, a, true, false)
            };
        }
        else if (transposeA && !transposeB) {
            return {
                a: () => matMul(b, dy, false, true),
                b: () => matMul(a, dy, false, false)
            };
        }
        else {
            return {
                a: () => matMul(b, dy, true, true),
                b: () => matMul(dy, a, true, true)
            };
        }
    }
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiQmF0Y2hNYXRNdWxfZ3JhZC5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uLy4uLy4uLy4uL3RmanMtY29yZS9zcmMvZ3JhZGllbnRzL0JhdGNoTWF0TXVsX2dyYWQudHMiXSwibmFtZXMiOltdLCJtYXBwaW5ncyI6IkFBQUE7Ozs7Ozs7Ozs7Ozs7OztHQWVHO0FBRUgsT0FBTyxFQUFDLFdBQVcsRUFBbUIsTUFBTSxpQkFBaUIsQ0FBQztBQUU5RCxPQUFPLEVBQUMsTUFBTSxFQUFDLE1BQU0sZ0JBQWdCLENBQUM7QUFHdEMsTUFBTSxDQUFDLE1BQU0scUJBQXFCLEdBQWU7SUFDL0MsVUFBVSxFQUFFLFdBQVc7SUFDdkIsWUFBWSxFQUFFLENBQUMsR0FBRyxFQUFFLEdBQUcsQ0FBQztJQUN4QixRQUFRLEVBQUUsQ0FBQyxFQUFVLEVBQUUsS0FBZSxFQUFFLEtBQW1CLEVBQUUsRUFBRTtRQUM3RCxNQUFNLENBQUMsQ0FBQyxFQUFFLENBQUMsQ0FBQyxHQUFHLEtBQUssQ0FBQztRQUVyQixNQUFNLEVBQUMsVUFBVSxFQUFFLFVBQVUsRUFBQyxHQUFHLEtBQW9DLENBQUM7UUFFdEUsSUFBSSxDQUFDLFVBQVUsSUFBSSxDQUFDLFVBQVUsRUFBRTtZQUM5QixPQUFPO2dCQUNMLENBQUMsRUFBRSxHQUFHLEVBQUUsQ0FBQyxNQUFNLENBQUMsRUFBRSxFQUFFLENBQUMsRUFBRSxLQUFLLEVBQUUsSUFBSSxDQUFDO2dCQUNuQyxDQUFDLEVBQUUsR0FBRyxFQUFFLENBQUMsTUFBTSxDQUFDLENBQUMsRUFBRSxFQUFFLEVBQUUsSUFBSSxFQUFFLEtBQUssQ0FBQzthQUNwQyxDQUFDO1NBQ0g7YUFBTSxJQUFJLENBQUMsVUFBVSxJQUFJLFVBQVUsRUFBRTtZQUNwQyxPQUFPO2dCQUNMLENBQUMsRUFBRSxHQUFHLEVBQUUsQ0FBQyxNQUFNLENBQUMsRUFBRSxFQUFFLENBQUMsRUFBRSxLQUFLLEVBQUUsS0FBSyxDQUFDO2dCQUNwQyxDQUFDLEVBQUUsR0FBRyxFQUFFLENBQUMsTUFBTSxDQUFDLEVBQUUsRUFBRSxDQUFDLEVBQUUsSUFBSSxFQUFFLEtBQUssQ0FBQzthQUNwQyxDQUFDO1NBQ0g7YUFBTSxJQUFJLFVBQVUsSUFBSSxDQUFDLFVBQVUsRUFBRTtZQUNwQyxPQUFPO2dCQUNMLENBQUMsRUFBRSxHQUFHLEVBQUUsQ0FBQyxNQUFNLENBQUMsQ0FBQyxFQUFFLEVBQUUsRUFBRSxLQUFLLEVBQUUsSUFBSSxDQUFDO2dCQUNuQyxDQUFDLEVBQUUsR0FBRyxFQUFFLENBQUMsTUFBTSxDQUFDLENBQUMsRUFBRSxFQUFFLEVBQUUsS0FBSyxFQUFFLEtBQUssQ0FBQzthQUNyQyxDQUFDO1NBQ0g7YUFBTTtZQUNMLE9BQU87Z0JBQ0wsQ0FBQyxFQUFFLEdBQUcsRUFBRSxDQUFDLE1BQU0sQ0FBQyxDQUFDLEVBQUUsRUFBRSxFQUFFLElBQUksRUFBRSxJQUFJLENBQUM7Z0JBQ2xDLENBQUMsRUFBRSxHQUFHLEVBQUUsQ0FBQyxNQUFNLENBQUMsRUFBRSxFQUFFLENBQUMsRUFBRSxJQUFJLEVBQUUsSUFBSSxDQUFDO2FBQ25DLENBQUM7U0FDSDtJQUNILENBQUM7Q0FDRixDQUFDIiwic291cmNlc0NvbnRlbnQiOlsiLyoqXG4gKiBAbGljZW5zZVxuICogQ29weXJpZ2h0IDIwMjAgR29vZ2xlIExMQy4gQWxsIFJpZ2h0cyBSZXNlcnZlZC5cbiAqIExpY2Vuc2VkIHVuZGVyIHRoZSBBcGFjaGUgTGljZW5zZSwgVmVyc2lvbiAyLjAgKHRoZSBcIkxpY2Vuc2VcIik7XG4gKiB5b3UgbWF5IG5vdCB1c2UgdGhpcyBmaWxlIGV4Y2VwdCBpbiBjb21wbGlhbmNlIHdpdGggdGhlIExpY2Vuc2UuXG4gKiBZb3UgbWF5IG9idGFpbiBhIGNvcHkgb2YgdGhlIExpY2Vuc2UgYXRcbiAqXG4gKiBodHRwOi8vd3d3LmFwYWNoZS5vcmcvbGljZW5zZXMvTElDRU5TRS0yLjBcbiAqXG4gKiBVbmxlc3MgcmVxdWlyZWQgYnkgYXBwbGljYWJsZSBsYXcgb3IgYWdyZWVkIHRvIGluIHdyaXRpbmcsIHNvZnR3YXJlXG4gKiBkaXN0cmlidXRlZCB1bmRlciB0aGUgTGljZW5zZSBpcyBkaXN0cmlidXRlZCBvbiBhbiBcIkFTIElTXCIgQkFTSVMsXG4gKiBXSVRIT1VUIFdBUlJBTlRJRVMgT1IgQ09ORElUSU9OUyBPRiBBTlkgS0lORCwgZWl0aGVyIGV4cHJlc3Mgb3IgaW1wbGllZC5cbiAqIFNlZSB0aGUgTGljZW5zZSBmb3IgdGhlIHNwZWNpZmljIGxhbmd1YWdlIGdvdmVybmluZyBwZXJtaXNzaW9ucyBhbmRcbiAqIGxpbWl0YXRpb25zIHVuZGVyIHRoZSBMaWNlbnNlLlxuICogPT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT1cbiAqL1xuXG5pbXBvcnQge0JhdGNoTWF0TXVsLCBCYXRjaE1hdE11bEF0dHJzfSBmcm9tICcuLi9rZXJuZWxfbmFtZXMnO1xuaW1wb3J0IHtHcmFkQ29uZmlnLCBOYW1lZEF0dHJNYXB9IGZyb20gJy4uL2tlcm5lbF9yZWdpc3RyeSc7XG5pbXBvcnQge21hdE11bH0gZnJvbSAnLi4vb3BzL21hdF9tdWwnO1xuaW1wb3J0IHtUZW5zb3J9IGZyb20gJy4uL3RlbnNvcic7XG5cbmV4cG9ydCBjb25zdCBiYXRjaE1hdE11bEdyYWRDb25maWc6IEdyYWRDb25maWcgPSB7XG4gIGtlcm5lbE5hbWU6IEJhdGNoTWF0TXVsLFxuICBpbnB1dHNUb1NhdmU6IFsnYScsICdiJ10sXG4gIGdyYWRGdW5jOiAoZHk6IFRlbnNvciwgc2F2ZWQ6IFRlbnNvcltdLCBhdHRyczogTmFtZWRBdHRyTWFwKSA9PiB7XG4gICAgY29uc3QgW2EsIGJdID0gc2F2ZWQ7XG5cbiAgICBjb25zdCB7dHJhbnNwb3NlQSwgdHJhbnNwb3NlQn0gPSBhdHRycyBhcyB1bmtub3duIGFzIEJhdGNoTWF0TXVsQXR0cnM7XG5cbiAgICBpZiAoIXRyYW5zcG9zZUEgJiYgIXRyYW5zcG9zZUIpIHtcbiAgICAgIHJldHVybiB7XG4gICAgICAgIGE6ICgpID0+IG1hdE11bChkeSwgYiwgZmFsc2UsIHRydWUpLFxuICAgICAgICBiOiAoKSA9PiBtYXRNdWwoYSwgZHksIHRydWUsIGZhbHNlKVxuICAgICAgfTtcbiAgICB9IGVsc2UgaWYgKCF0cmFuc3Bvc2VBICYmIHRyYW5zcG9zZUIpIHtcbiAgICAgIHJldHVybiB7XG4gICAgICAgIGE6ICgpID0+IG1hdE11bChkeSwgYiwgZmFsc2UsIGZhbHNlKSxcbiAgICAgICAgYjogKCkgPT4gbWF0TXVsKGR5LCBhLCB0cnVlLCBmYWxzZSlcbiAgICAgIH07XG4gICAgfSBlbHNlIGlmICh0cmFuc3Bvc2VBICYmICF0cmFuc3Bvc2VCKSB7XG4gICAgICByZXR1cm4ge1xuICAgICAgICBhOiAoKSA9PiBtYXRNdWwoYiwgZHksIGZhbHNlLCB0cnVlKSxcbiAgICAgICAgYjogKCkgPT4gbWF0TXVsKGEsIGR5LCBmYWxzZSwgZmFsc2UpXG4gICAgICB9O1xuICAgIH0gZWxzZSB7XG4gICAgICByZXR1cm4ge1xuICAgICAgICBhOiAoKSA9PiBtYXRNdWwoYiwgZHksIHRydWUsIHRydWUpLFxuICAgICAgICBiOiAoKSA9PiBtYXRNdWwoZHksIGEsIHRydWUsIHRydWUpXG4gICAgICB9O1xuICAgIH1cbiAgfVxufTtcbiJdfQ==