UNPKG

@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

55 lines 7.77 kB
/** * @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 { Reverse, util } from '@tensorflow/tfjs-core'; import { identity } from './Identity'; import { reshape } from './Reshape'; let wasmReverse; function setup(backend) { wasmReverse = backend.wasm.cwrap(Reverse, null, [ 'number', 'array', 'number', 'array', 'number', 'number' // out_id ]); } export function reverse(args) { const { inputs, backend, attrs } = args; const { x } = inputs; const { dims } = attrs; const axes = util.parseAxisParam(dims, x.shape); if (x.shape.length === 0) { return identity({ inputs: { x }, backend }); } const out = backend.makeOutput(x.shape, x.dtype); const xId = backend.dataIdMap.get(x.dataId).id; const outId = backend.dataIdMap.get(out.dataId).id; const axesBytes = new Uint8Array(new Int32Array(axes).buffer); const outShapeBytes = new Uint8Array(new Int32Array(x.shape).buffer); wasmReverse(xId, axesBytes, axes.length, outShapeBytes, x.shape.length, outId); const reshaped = reshape({ inputs: { x: out }, attrs: { shape: x.shape }, backend }); backend.disposeData(out.dataId); return reshaped; } export const reverseConfig = { kernelName: Reverse, backendName: 'wasm', kernelFunc: reverse, setupFunc: setup }; //# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiUmV2ZXJzZS5qcyIsInNvdXJjZVJvb3QiOiIiLCJzb3VyY2VzIjpbIi4uLy4uLy4uLy4uLy4uLy4uL3RmanMtYmFja2VuZC13YXNtL3NyYy9rZXJuZWxzL1JldmVyc2UudHMiXSwibmFtZXMiOltdLCJtYXBwaW5ncyI6IkFBQUE7Ozs7Ozs7Ozs7Ozs7OztHQWVHO0FBRUgsT0FBTyxFQUEyQixPQUFPLEVBQTJDLElBQUksRUFBQyxNQUFNLHVCQUF1QixDQUFDO0FBSXZILE9BQU8sRUFBQyxRQUFRLEVBQUMsTUFBTSxZQUFZLENBQUM7QUFDcEMsT0FBTyxFQUFDLE9BQU8sRUFBQyxNQUFNLFdBQVcsQ0FBQztBQUVsQyxJQUFJLFdBRThDLENBQUM7QUFFbkQsU0FBUyxLQUFLLENBQUMsT0FBb0I7SUFDakMsV0FBVyxHQUFHLE9BQU8sQ0FBQyxJQUFJLENBQUMsS0FBSyxDQUFDLE9BQU8sRUFBRSxJQUFJLEVBQUU7UUFDOUMsUUFBUTtRQUNSLE9BQU87UUFDUCxRQUFRO1FBQ1IsT0FBTztRQUNQLFFBQVE7UUFDUixRQUFRLENBQUcsU0FBUztLQUNyQixDQUFDLENBQUM7QUFDTCxDQUFDO0FBRUQsTUFBTSxVQUFVLE9BQU8sQ0FDbkIsSUFBd0U7SUFFMUUsTUFBTSxFQUFDLE1BQU0sRUFBRSxPQUFPLEVBQUUsS0FBSyxFQUFDLEdBQUcsSUFBSSxDQUFDO0lBQ3RDLE1BQU0sRUFBQyxDQUFDLEVBQUMsR0FBRyxNQUFNLENBQUM7SUFDbkIsTUFBTSxFQUFDLElBQUksRUFBQyxHQUFHLEtBQUssQ0FBQztJQUVyQixNQUFNLElBQUksR0FBRyxJQUFJLENBQUMsY0FBYyxDQUFDLElBQUksRUFBRSxDQUFDLENBQUMsS0FBSyxDQUFDLENBQUM7SUFFaEQsSUFBSSxDQUFDLENBQUMsS0FBSyxDQUFDLE1BQU0sS0FBSyxDQUFDLEVBQUU7UUFDeEIsT0FBTyxRQUFRLENBQUMsRUFBQyxNQUFNLEVBQUUsRUFBQyxDQUFDLEVBQUMsRUFBRSxPQUFPLEVBQUMsQ0FBQyxDQUFDO0tBQ3pDO0lBRUQsTUFBTSxHQUFHLEdBQUcsT0FBTyxDQUFDLFVBQVUsQ0FBQyxDQUFDLENBQUMsS0FBSyxFQUFFLENBQUMsQ0FBQyxLQUFLLENBQUMsQ0FBQztJQUNqRCxNQUFNLEdBQUcsR0FBRyxPQUFPLENBQUMsU0FBUyxDQUFDLEdBQUcsQ0FBQyxDQUFDLENBQUMsTUFBTSxDQUFDLENBQUMsRUFBRSxDQUFDO0lBQy9DLE1BQU0sS0FBSyxHQUFHLE9BQU8sQ0FBQyxTQUFTLENBQUMsR0FBRyxDQUFDLEdBQUcsQ0FBQyxNQUFNLENBQUMsQ0FBQyxFQUFFLENBQUM7SUFFbkQsTUFBTSxTQUFTLEdBQUcsSUFBSSxVQUFVLENBQUMsSUFBSSxVQUFVLENBQUMsSUFBSSxDQUFDLENBQUMsTUFBTSxDQUFDLENBQUM7SUFDOUQsTUFBTSxhQUFhLEdBQUcsSUFBSSxVQUFVLENBQUMsSUFBSSxVQUFVLENBQUMsQ0FBQyxDQUFDLEtBQUssQ0FBQyxDQUFDLE1BQU0sQ0FBQyxDQUFDO0lBRXJFLFdBQVcsQ0FDUCxHQUFHLEVBQUUsU0FBUyxFQUFFLElBQUksQ0FBQyxNQUFNLEVBQUUsYUFBYSxFQUFFLENBQUMsQ0FBQyxLQUFLLENBQUMsTUFBTSxFQUFFLEtBQUssQ0FBQyxDQUFDO0lBRXZFLE1BQU0sUUFBUSxHQUNWLE9BQU8sQ0FBQyxFQUFDLE1BQU0sRUFBRSxFQUFDLENBQUMsRUFBRSxHQUFHLEVBQUMsRUFBRSxLQUFLLEVBQUUsRUFBQyxLQUFLLEVBQUUsQ0FBQyxDQUFDLEtBQUssRUFBQyxFQUFFLE9BQU8sRUFBQyxDQUFDLENBQUM7SUFFbEUsT0FBTyxDQUFDLFdBQVcsQ0FBQyxHQUFHLENBQUMsTUFBTSxDQUFDLENBQUM7SUFDaEMsT0FBTyxRQUFRLENBQUM7QUFDbEIsQ0FBQztBQUVELE1BQU0sQ0FBQyxNQUFNLGFBQWEsR0FBaUI7SUFDekMsVUFBVSxFQUFFLE9BQU87SUFDbkIsV0FBVyxFQUFFLE1BQU07SUFDbkIsVUFBVSxFQUFFLE9BQWdDO0lBQzVDLFNBQVMsRUFBRSxLQUFLO0NBQ2pCLENBQUMiLCJzb3VyY2VzQ29udGVudCI6WyIvKipcbiAqIEBsaWNlbnNlXG4gKiBDb3B5cmlnaHQgMjAyMCBHb29nbGUgTExDLiBBbGwgUmlnaHRzIFJlc2VydmVkLlxuICogTGljZW5zZWQgdW5kZXIgdGhlIEFwYWNoZSBMaWNlbnNlLCBWZXJzaW9uIDIuMCAodGhlIFwiTGljZW5zZVwiKTtcbiAqIHlvdSBtYXkgbm90IHVzZSB0aGlzIGZpbGUgZXhjZXB0IGluIGNvbXBsaWFuY2Ugd2l0aCB0aGUgTGljZW5zZS5cbiAqIFlvdSBtYXkgb2J0YWluIGEgY29weSBvZiB0aGUgTGljZW5zZSBhdFxuICpcbiAqIGh0dHA6Ly93d3cuYXBhY2hlLm9yZy9saWNlbnNlcy9MSUNFTlNFLTIuMFxuICpcbiAqIFVubGVzcyByZXF1aXJlZCBieSBhcHBsaWNhYmxlIGxhdyBvciBhZ3JlZWQgdG8gaW4gd3JpdGluZywgc29mdHdhcmVcbiAqIGRpc3RyaWJ1dGVkIHVuZGVyIHRoZSBMaWNlbnNlIGlzIGRpc3RyaWJ1dGVkIG9uIGFuIFwiQVMgSVNcIiBCQVNJUyxcbiAqIFdJVEhPVVQgV0FSUkFOVElFUyBPUiBDT05ESVRJT05TIE9GIEFOWSBLSU5ELCBlaXRoZXIgZXhwcmVzcyBvciBpbXBsaWVkLlxuICogU2VlIHRoZSBMaWNlbnNlIGZvciB0aGUgc3BlY2lmaWMgbGFuZ3VhZ2UgZ292ZXJuaW5nIHBlcm1pc3Npb25zIGFuZFxuICogbGltaXRhdGlvbnMgdW5kZXIgdGhlIExpY2Vuc2UuXG4gKiA9PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PVxuICovXG5cbmltcG9ydCB7S2VybmVsQ29uZmlnLCBLZXJuZWxGdW5jLCBSZXZlcnNlLCBSZXZlcnNlQXR0cnMsIFJldmVyc2VJbnB1dHMsIFRlbnNvckluZm8sIHV0aWx9IGZyb20gJ0B0ZW5zb3JmbG93L3RmanMtY29yZSc7XG5cbmltcG9ydCB7QmFja2VuZFdhc219IGZyb20gJy4uL2JhY2tlbmRfd2FzbSc7XG5cbmltcG9ydCB7aWRlbnRpdHl9IGZyb20gJy4vSWRlbnRpdHknO1xuaW1wb3J0IHtyZXNoYXBlfSBmcm9tICcuL1Jlc2hhcGUnO1xuXG5sZXQgd2FzbVJldmVyc2U6IChcbiAgICB4SWQ6IG51bWJlciwgYXhlczogVWludDhBcnJheSwgYXhlc0xlbmd0aDogbnVtYmVyLCBvdXRTaGFwZTogVWludDhBcnJheSxcbiAgICBvdXRTaGFwZUxlbmd0aDogbnVtYmVyLCBvdXRJZDogbnVtYmVyKSA9PiB2b2lkO1xuXG5mdW5jdGlvbiBzZXR1cChiYWNrZW5kOiBCYWNrZW5kV2FzbSkge1xuICB3YXNtUmV2ZXJzZSA9IGJhY2tlbmQud2FzbS5jd3JhcChSZXZlcnNlLCBudWxsLCBbXG4gICAgJ251bWJlcicsICAvLyB4X2lkXG4gICAgJ2FycmF5JywgICAvLyBheGVzXG4gICAgJ251bWJlcicsICAvLyBheGVzX2xlbmd0aFxuICAgICdhcnJheScsICAgLy8gb3V0X3NoYXBlXG4gICAgJ251bWJlcicsICAvLyBvdXRfc2hhcGVfbGVuZ3RoXG4gICAgJ251bWJlcicgICAvLyBvdXRfaWRcbiAgXSk7XG59XG5cbmV4cG9ydCBmdW5jdGlvbiByZXZlcnNlKFxuICAgIGFyZ3M6IHtpbnB1dHM6IFJldmVyc2VJbnB1dHMsIGJhY2tlbmQ6IEJhY2tlbmRXYXNtLCBhdHRyczogUmV2ZXJzZUF0dHJzfSk6XG4gICAgVGVuc29ySW5mbyB7XG4gIGNvbnN0IHtpbnB1dHMsIGJhY2tlbmQsIGF0dHJzfSA9IGFyZ3M7XG4gIGNvbnN0IHt4fSA9IGlucHV0cztcbiAgY29uc3Qge2RpbXN9ID0gYXR0cnM7XG5cbiAgY29uc3QgYXhlcyA9IHV0aWwucGFyc2VBeGlzUGFyYW0oZGltcywgeC5zaGFwZSk7XG5cbiAgaWYgKHguc2hhcGUubGVuZ3RoID09PSAwKSB7XG4gICAgcmV0dXJuIGlkZW50aXR5KHtpbnB1dHM6IHt4fSwgYmFja2VuZH0pO1xuICB9XG5cbiAgY29uc3Qgb3V0ID0gYmFja2VuZC5tYWtlT3V0cHV0KHguc2hhcGUsIHguZHR5cGUpO1xuICBjb25zdCB4SWQgPSBiYWNrZW5kLmRhdGFJZE1hcC5nZXQoeC5kYXRhSWQpLmlkO1xuICBjb25zdCBvdXRJZCA9IGJhY2tlbmQuZGF0YUlkTWFwLmdldChvdXQuZGF0YUlkKS5pZDtcblxuICBjb25zdCBheGVzQnl0ZXMgPSBuZXcgVWludDhBcnJheShuZXcgSW50MzJBcnJheShheGVzKS5idWZmZXIpO1xuICBjb25zdCBvdXRTaGFwZUJ5dGVzID0gbmV3IFVpbnQ4QXJyYXkobmV3IEludDMyQXJyYXkoeC5zaGFwZSkuYnVmZmVyKTtcblxuICB3YXNtUmV2ZXJzZShcbiAgICAgIHhJZCwgYXhlc0J5dGVzLCBheGVzLmxlbmd0aCwgb3V0U2hhcGVCeXRlcywgeC5zaGFwZS5sZW5ndGgsIG91dElkKTtcblxuICBjb25zdCByZXNoYXBlZCA9XG4gICAgICByZXNoYXBlKHtpbnB1dHM6IHt4OiBvdXR9LCBhdHRyczoge3NoYXBlOiB4LnNoYXBlfSwgYmFja2VuZH0pO1xuXG4gIGJhY2tlbmQuZGlzcG9zZURhdGEob3V0LmRhdGFJZCk7XG4gIHJldHVybiByZXNoYXBlZDtcbn1cblxuZXhwb3J0IGNvbnN0IHJldmVyc2VDb25maWc6IEtlcm5lbENvbmZpZyA9IHtcbiAga2VybmVsTmFtZTogUmV2ZXJzZSxcbiAgYmFja2VuZE5hbWU6ICd3YXNtJyxcbiAga2VybmVsRnVuYzogcmV2ZXJzZSBhcyB1bmtub3duIGFzIEtlcm5lbEZ1bmMsXG4gIHNldHVwRnVuYzogc2V0dXBcbn07XG4iXX0=