@tensorflow/tfjs-backend-wasm
Version:
This package adds a WebAssembly backend to TensorFlow.js. It currently supports the following models from our [models](https://github.com/tensorflow/tfjs-models) repo: - BlazeFace - BodyPix - CocoSSD - Face landmarks detection - HandPose - KNN classifier
49 lines • 8.29 kB
JavaScript
/**
* @license
* Copyright 2021 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 { backend_util, BatchToSpaceND } from '@tensorflow/tfjs-core';
import { reshape } from './Reshape';
import { slice } from './Slice';
import { transpose } from './Transpose';
function batchToSpaceND(args) {
const { inputs, backend, attrs } = args;
const { x } = inputs;
const { blockShape, crops } = attrs;
const prod = blockShape.reduce((a, b) => a * b);
const reshaped = backend_util.getReshaped(x.shape, blockShape, prod);
const permuted = backend_util.getPermuted(reshaped.length, blockShape.length);
const reshapedPermuted = backend_util.getReshapedPermuted(x.shape, blockShape, prod);
const sliceBeginCoords = backend_util.getSliceBeginCoords(crops, blockShape.length);
const sliceSize = backend_util.getSliceSize(reshapedPermuted, crops, blockShape.length);
const xReshaped = reshape({ inputs: { x }, backend, attrs: { shape: reshaped } });
const xTransposed = transpose({ inputs: { x: xReshaped }, backend, attrs: { perm: permuted } });
const xTransposedReshaped = reshape({ inputs: { x: xTransposed }, backend, attrs: { shape: reshapedPermuted } });
const result = slice({
inputs: { x: xTransposedReshaped },
backend,
attrs: { begin: sliceBeginCoords, size: sliceSize }
});
backend.disposeData(xReshaped.dataId);
backend.disposeData(xTransposed.dataId);
backend.disposeData(xTransposedReshaped.dataId);
return result;
}
export const batchToSpaceNDConfig = {
kernelName: BatchToSpaceND,
backendName: 'wasm',
kernelFunc: batchToSpaceND
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiQmF0Y2hUb1NwYWNlTkQuanMiLCJzb3VyY2VSb290IjoiIiwic291cmNlcyI6WyIuLi8uLi8uLi8uLi8uLi8uLi90ZmpzLWJhY2tlbmQtd2FzbS9zcmMva2VybmVscy9CYXRjaFRvU3BhY2VORC50cyJdLCJuYW1lcyI6W10sIm1hcHBpbmdzIjoiQUFBQTs7Ozs7Ozs7Ozs7Ozs7O0dBZUc7QUFFSCxPQUFPLEVBQUMsWUFBWSxFQUFFLGNBQWMsRUFBc0UsTUFBTSx1QkFBdUIsQ0FBQztBQUl4SSxPQUFPLEVBQUMsT0FBTyxFQUFDLE1BQU0sV0FBVyxDQUFDO0FBQ2xDLE9BQU8sRUFBQyxLQUFLLEVBQUMsTUFBTSxTQUFTLENBQUM7QUFDOUIsT0FBTyxFQUFDLFNBQVMsRUFBQyxNQUFNLGFBQWEsQ0FBQztBQUV0QyxTQUFTLGNBQWMsQ0FBQyxJQUl2QjtJQUNDLE1BQU0sRUFBQyxNQUFNLEVBQUUsT0FBTyxFQUFFLEtBQUssRUFBQyxHQUFHLElBQUksQ0FBQztJQUN0QyxNQUFNLEVBQUMsQ0FBQyxFQUFDLEdBQUcsTUFBTSxDQUFDO0lBQ25CLE1BQU0sRUFBQyxVQUFVLEVBQUUsS0FBSyxFQUFDLEdBQUcsS0FBSyxDQUFDO0lBRWxDLE1BQU0sSUFBSSxHQUFHLFVBQVUsQ0FBQyxNQUFNLENBQUMsQ0FBQyxDQUFDLEVBQUUsQ0FBQyxFQUFFLEVBQUUsQ0FBQyxDQUFDLEdBQUcsQ0FBQyxDQUFDLENBQUM7SUFFaEQsTUFBTSxRQUFRLEdBQUcsWUFBWSxDQUFDLFdBQVcsQ0FBQyxDQUFDLENBQUMsS0FBSyxFQUFFLFVBQVUsRUFBRSxJQUFJLENBQUMsQ0FBQztJQUNyRSxNQUFNLFFBQVEsR0FBRyxZQUFZLENBQUMsV0FBVyxDQUFDLFFBQVEsQ0FBQyxNQUFNLEVBQUUsVUFBVSxDQUFDLE1BQU0sQ0FBQyxDQUFDO0lBQzlFLE1BQU0sZ0JBQWdCLEdBQ2xCLFlBQVksQ0FBQyxtQkFBbUIsQ0FBQyxDQUFDLENBQUMsS0FBSyxFQUFFLFVBQVUsRUFBRSxJQUFJLENBQUMsQ0FBQztJQUNoRSxNQUFNLGdCQUFnQixHQUNsQixZQUFZLENBQUMsbUJBQW1CLENBQUMsS0FBSyxFQUFFLFVBQVUsQ0FBQyxNQUFNLENBQUMsQ0FBQztJQUMvRCxNQUFNLFNBQVMsR0FDWCxZQUFZLENBQUMsWUFBWSxDQUFDLGdCQUFnQixFQUFFLEtBQUssRUFBRSxVQUFVLENBQUMsTUFBTSxDQUFDLENBQUM7SUFFMUUsTUFBTSxTQUFTLEdBQUcsT0FBTyxDQUFDLEVBQUMsTUFBTSxFQUFFLEVBQUMsQ0FBQyxFQUFDLEVBQUUsT0FBTyxFQUFFLEtBQUssRUFBRSxFQUFDLEtBQUssRUFBRSxRQUFRLEVBQUMsRUFBQyxDQUFDLENBQUM7SUFDNUUsTUFBTSxXQUFXLEdBQ2IsU0FBUyxDQUFDLEVBQUMsTUFBTSxFQUFFLEVBQUMsQ0FBQyxFQUFFLFNBQVMsRUFBQyxFQUFFLE9BQU8sRUFBRSxLQUFLLEVBQUUsRUFBQyxJQUFJLEVBQUUsUUFBUSxFQUFDLEVBQUMsQ0FBQyxDQUFDO0lBQzFFLE1BQU0sbUJBQW1CLEdBQUcsT0FBTyxDQUMvQixFQUFDLE1BQU0sRUFBRSxFQUFDLENBQUMsRUFBRSxXQUFXLEVBQUMsRUFBRSxPQUFPLEVBQUUsS0FBSyxFQUFFLEVBQUMsS0FBSyxFQUFFLGdCQUFnQixFQUFDLEVBQUMsQ0FBQyxDQUFDO0lBQzNFLE1BQU0sTUFBTSxHQUFHLEtBQUssQ0FBQztRQUNuQixNQUFNLEVBQUUsRUFBQyxDQUFDLEVBQUUsbUJBQW1CLEVBQUM7UUFDaEMsT0FBTztRQUNQLEtBQUssRUFBRSxFQUFDLEtBQUssRUFBRSxnQkFBZ0IsRUFBRSxJQUFJLEVBQUUsU0FBUyxFQUFDO0tBQ2xELENBQUMsQ0FBQztJQUVILE9BQU8sQ0FBQyxXQUFXLENBQUMsU0FBUyxDQUFDLE1BQU0sQ0FBQyxDQUFDO0lBQ3RDLE9BQU8sQ0FBQyxXQUFXLENBQUMsV0FBVyxDQUFDLE1BQU0sQ0FBQyxDQUFDO0lBQ3hDLE9BQU8sQ0FBQyxXQUFXLENBQUMsbUJBQW1CLENBQUMsTUFBTSxDQUFDLENBQUM7SUFFaEQsT0FBTyxNQUFNLENBQUM7QUFDaEIsQ0FBQztBQUVELE1BQU0sQ0FBQyxNQUFNLG9CQUFvQixHQUFpQjtJQUNoRCxVQUFVLEVBQUUsY0FBYztJQUMxQixXQUFXLEVBQUUsTUFBTTtJQUNuQixVQUFVLEVBQUUsY0FBdUM7Q0FDcEQsQ0FBQyIsInNvdXJjZXNDb250ZW50IjpbIi8qKlxuICogQGxpY2Vuc2VcbiAqIENvcHlyaWdodCAyMDIxIEdvb2dsZSBMTEMuIEFsbCBSaWdodHMgUmVzZXJ2ZWQuXG4gKiBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgXCJMaWNlbnNlXCIpO1xuICogeW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLlxuICogWW91IG1heSBvYnRhaW4gYSBjb3B5IG9mIHRoZSBMaWNlbnNlIGF0XG4gKlxuICogaHR0cDovL3d3dy5hcGFjaGUub3JnL2xpY2Vuc2VzL0xJQ0VOU0UtMi4wXG4gKlxuICogVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZVxuICogZGlzdHJpYnV0ZWQgdW5kZXIgdGhlIExpY2Vuc2UgaXMgZGlzdHJpYnV0ZWQgb24gYW4gXCJBUyBJU1wiIEJBU0lTLFxuICogV0lUSE9VVCBXQVJSQU5USUVTIE9SIENPTkRJVElPTlMgT0YgQU5ZIEtJTkQsIGVpdGhlciBleHByZXNzIG9yIGltcGxpZWQuXG4gKiBTZWUgdGhlIExpY2Vuc2UgZm9yIHRoZSBzcGVjaWZpYyBsYW5ndWFnZSBnb3Zlcm5pbmcgcGVybWlzc2lvbnMgYW5kXG4gKiBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS5cbiAqID09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09XG4gKi9cblxuaW1wb3J0IHtiYWNrZW5kX3V0aWwsIEJhdGNoVG9TcGFjZU5ELCBCYXRjaFRvU3BhY2VOREF0dHJzLCBCYXRjaFRvU3BhY2VORElucHV0cywgS2VybmVsQ29uZmlnLCBLZXJuZWxGdW5jfSBmcm9tICdAdGVuc29yZmxvdy90ZmpzLWNvcmUnO1xuXG5pbXBvcnQge0JhY2tlbmRXYXNtfSBmcm9tICcuLi9iYWNrZW5kX3dhc20nO1xuXG5pbXBvcnQge3Jlc2hhcGV9IGZyb20gJy4vUmVzaGFwZSc7XG5pbXBvcnQge3NsaWNlfSBmcm9tICcuL1NsaWNlJztcbmltcG9ydCB7dHJhbnNwb3NlfSBmcm9tICcuL1RyYW5zcG9zZSc7XG5cbmZ1bmN0aW9uIGJhdGNoVG9TcGFjZU5EKGFyZ3M6IHtcbiAgaW5wdXRzOiBCYXRjaFRvU3BhY2VORElucHV0cyxcbiAgYmFja2VuZDogQmFja2VuZFdhc20sXG4gIGF0dHJzOiBCYXRjaFRvU3BhY2VOREF0dHJzXG59KSB7XG4gIGNvbnN0IHtpbnB1dHMsIGJhY2tlbmQsIGF0dHJzfSA9IGFyZ3M7XG4gIGNvbnN0IHt4fSA9IGlucHV0cztcbiAgY29uc3Qge2Jsb2NrU2hhcGUsIGNyb3BzfSA9IGF0dHJzO1xuXG4gIGNvbnN0IHByb2QgPSBibG9ja1NoYXBlLnJlZHVjZSgoYSwgYikgPT4gYSAqIGIpO1xuXG4gIGNvbnN0IHJlc2hhcGVkID0gYmFja2VuZF91dGlsLmdldFJlc2hhcGVkKHguc2hhcGUsIGJsb2NrU2hhcGUsIHByb2QpO1xuICBjb25zdCBwZXJtdXRlZCA9IGJhY2tlbmRfdXRpbC5nZXRQZXJtdXRlZChyZXNoYXBlZC5sZW5ndGgsIGJsb2NrU2hhcGUubGVuZ3RoKTtcbiAgY29uc3QgcmVzaGFwZWRQZXJtdXRlZCA9XG4gICAgICBiYWNrZW5kX3V0aWwuZ2V0UmVzaGFwZWRQZXJtdXRlZCh4LnNoYXBlLCBibG9ja1NoYXBlLCBwcm9kKTtcbiAgY29uc3Qgc2xpY2VCZWdpbkNvb3JkcyA9XG4gICAgICBiYWNrZW5kX3V0aWwuZ2V0U2xpY2VCZWdpbkNvb3Jkcyhjcm9wcywgYmxvY2tTaGFwZS5sZW5ndGgpO1xuICBjb25zdCBzbGljZVNpemUgPVxuICAgICAgYmFja2VuZF91dGlsLmdldFNsaWNlU2l6ZShyZXNoYXBlZFBlcm11dGVkLCBjcm9wcywgYmxvY2tTaGFwZS5sZW5ndGgpO1xuXG4gIGNvbnN0IHhSZXNoYXBlZCA9IHJlc2hhcGUoe2lucHV0czoge3h9LCBiYWNrZW5kLCBhdHRyczoge3NoYXBlOiByZXNoYXBlZH19KTtcbiAgY29uc3QgeFRyYW5zcG9zZWQgPVxuICAgICAgdHJhbnNwb3NlKHtpbnB1dHM6IHt4OiB4UmVzaGFwZWR9LCBiYWNrZW5kLCBhdHRyczoge3Blcm06IHBlcm11dGVkfX0pO1xuICBjb25zdCB4VHJhbnNwb3NlZFJlc2hhcGVkID0gcmVzaGFwZShcbiAgICAgIHtpbnB1dHM6IHt4OiB4VHJhbnNwb3NlZH0sIGJhY2tlbmQsIGF0dHJzOiB7c2hhcGU6IHJlc2hhcGVkUGVybXV0ZWR9fSk7XG4gIGNvbnN0IHJlc3VsdCA9IHNsaWNlKHtcbiAgICBpbnB1dHM6IHt4OiB4VHJhbnNwb3NlZFJlc2hhcGVkfSxcbiAgICBiYWNrZW5kLFxuICAgIGF0dHJzOiB7YmVnaW46IHNsaWNlQmVnaW5Db29yZHMsIHNpemU6IHNsaWNlU2l6ZX1cbiAgfSk7XG5cbiAgYmFja2VuZC5kaXNwb3NlRGF0YSh4UmVzaGFwZWQuZGF0YUlkKTtcbiAgYmFja2VuZC5kaXNwb3NlRGF0YSh4VHJhbnNwb3NlZC5kYXRhSWQpO1xuICBiYWNrZW5kLmRpc3Bvc2VEYXRhKHhUcmFuc3Bvc2VkUmVzaGFwZWQuZGF0YUlkKTtcblxuICByZXR1cm4gcmVzdWx0O1xufVxuXG5leHBvcnQgY29uc3QgYmF0Y2hUb1NwYWNlTkRDb25maWc6IEtlcm5lbENvbmZpZyA9IHtcbiAga2VybmVsTmFtZTogQmF0Y2hUb1NwYWNlTkQsXG4gIGJhY2tlbmROYW1lOiAnd2FzbScsXG4gIGtlcm5lbEZ1bmM6IGJhdGNoVG9TcGFjZU5EIGFzIHVua25vd24gYXMgS2VybmVsRnVuY1xufTtcbiJdfQ==