@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
57 lines • 7.59 kB
JavaScript
/**
* @license
* Copyright 2023 Google LLC.
* 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 { SearchSorted } from '@tensorflow/tfjs-core';
import { CppDType } from './types';
let wasmSearchSorted;
function setup(backend) {
wasmSearchSorted = backend.wasm.cwrap(SearchSorted, null /* void */, [
'number',
'number',
'number',
'number',
'number',
'number',
'bool',
'number', // outId
]);
}
function searchSorted(args) {
const { inputs, backend, attrs } = args;
const { sortedSequence, values } = inputs;
const { side } = attrs;
if (sortedSequence.dtype !== values.dtype) {
throw new Error(`SearchSorted error: sorted_sequence must have the same dtype as values. Got ${sortedSequence.dtype} and ${values.dtype}`);
}
const out = backend.makeOutput(values.shape, 'int32');
function tensorId(x) {
return backend.dataIdMap.get(x.dataId).id;
}
wasmSearchSorted(tensorId(sortedSequence), tensorId(values),
/*batchSize=*/ sortedSequence.shape[0],
/*sequenceSize=*/ sortedSequence.shape[1],
/*valuesSize=*/ values.shape[1],
/*dtype=*/ CppDType[sortedSequence.dtype],
/*isSideLeft=*/ side === 'left', tensorId(out));
return out;
}
export const searchSortedConfig = {
kernelName: SearchSorted,
backendName: 'wasm',
setupFunc: setup,
kernelFunc: searchSorted
};
//# sourceMappingURL=data:application/json;base64,eyJ2ZXJzaW9uIjozLCJmaWxlIjoiU2VhcmNoU29ydGVkLmpzIiwic291cmNlUm9vdCI6IiIsInNvdXJjZXMiOlsiLi4vLi4vLi4vLi4vLi4vLi4vdGZqcy1iYWNrZW5kLXdhc20vc3JjL2tlcm5lbHMvU2VhcmNoU29ydGVkLnRzIl0sIm5hbWVzIjpbXSwibWFwcGluZ3MiOiJBQUFBOzs7Ozs7Ozs7Ozs7Ozs7R0FlRztBQUVILE9BQU8sRUFBMkIsWUFBWSxFQUFvRCxNQUFNLHVCQUF1QixDQUFDO0FBR2hJLE9BQU8sRUFBQyxRQUFRLEVBQUMsTUFBTSxTQUFTLENBQUM7QUFFakMsSUFBSSxnQkFHMkMsQ0FBQztBQUVoRCxTQUFTLEtBQUssQ0FBQyxPQUFvQjtJQUNqQyxnQkFBZ0IsR0FBRyxPQUFPLENBQUMsSUFBSSxDQUFDLEtBQUssQ0FBQyxZQUFZLEVBQUUsSUFBSSxDQUFDLFVBQVUsRUFBRTtRQUNuRSxRQUFRO1FBQ1IsUUFBUTtRQUNSLFFBQVE7UUFDUixRQUFRO1FBQ1IsUUFBUTtRQUNSLFFBQVE7UUFDUixNQUFNO1FBQ04sUUFBUSxFQUFHLFFBQVE7S0FDcEIsQ0FBQyxDQUFDO0FBQ0wsQ0FBQztBQUVELFNBQVMsWUFBWSxDQUFDLElBSXJCO0lBQ0MsTUFBTSxFQUFDLE1BQU0sRUFBRSxPQUFPLEVBQUUsS0FBSyxFQUFDLEdBQUcsSUFBSSxDQUFDO0lBQ3RDLE1BQU0sRUFBQyxjQUFjLEVBQUUsTUFBTSxFQUFDLEdBQUcsTUFBTSxDQUFDO0lBQ3hDLE1BQU0sRUFBQyxJQUFJLEVBQUMsR0FBRyxLQUFLLENBQUM7SUFFckIsSUFBSSxjQUFjLENBQUMsS0FBSyxLQUFLLE1BQU0sQ0FBQyxLQUFLLEVBQUU7UUFDekMsTUFBTSxJQUFJLEtBQUssQ0FDWCwrRUFDSSxjQUFjLENBQUMsS0FBSyxRQUFRLE1BQU0sQ0FBQyxLQUFLLEVBQUUsQ0FBQyxDQUFDO0tBQ3JEO0lBRUQsTUFBTSxHQUFHLEdBQUcsT0FBTyxDQUFDLFVBQVUsQ0FBQyxNQUFNLENBQUMsS0FBSyxFQUFFLE9BQU8sQ0FBQyxDQUFDO0lBRXRELFNBQVMsUUFBUSxDQUFDLENBQWE7UUFDN0IsT0FBTyxPQUFPLENBQUMsU0FBUyxDQUFDLEdBQUcsQ0FBQyxDQUFDLENBQUMsTUFBTSxDQUFDLENBQUMsRUFBRyxDQUFDO0lBQzdDLENBQUM7SUFDRCxnQkFBZ0IsQ0FDWixRQUFRLENBQUMsY0FBYyxDQUFDLEVBQ3hCLFFBQVEsQ0FBQyxNQUFNLENBQUM7SUFDaEIsY0FBYyxDQUFBLGNBQWMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDO0lBQ3JDLGlCQUFpQixDQUFBLGNBQWMsQ0FBQyxLQUFLLENBQUMsQ0FBQyxDQUFDO0lBQ3hDLGVBQWUsQ0FBQSxNQUFNLENBQUMsS0FBSyxDQUFDLENBQUMsQ0FBQztJQUM5QixVQUFVLENBQUEsUUFBUSxDQUFDLGNBQWMsQ0FBQyxLQUFLLENBQUM7SUFDeEMsZUFBZSxDQUFBLElBQUksS0FBSyxNQUFNLEVBQzlCLFFBQVEsQ0FBQyxHQUFHLENBQUMsQ0FDaEIsQ0FBQztJQUVGLE9BQU8sR0FBRyxDQUFDO0FBQ2IsQ0FBQztBQUVELE1BQU0sQ0FBQyxNQUFNLGtCQUFrQixHQUFpQjtJQUM5QyxVQUFVLEVBQUUsWUFBWTtJQUN4QixXQUFXLEVBQUUsTUFBTTtJQUNuQixTQUFTLEVBQUUsS0FBSztJQUNoQixVQUFVLEVBQUUsWUFBcUM7Q0FDbEQsQ0FBQyIsInNvdXJjZXNDb250ZW50IjpbIi8qKlxuICogQGxpY2Vuc2VcbiAqIENvcHlyaWdodCAyMDIzIEdvb2dsZSBMTEMuXG4gKiBMaWNlbnNlZCB1bmRlciB0aGUgQXBhY2hlIExpY2Vuc2UsIFZlcnNpb24gMi4wICh0aGUgXCJMaWNlbnNlXCIpO1xuICogeW91IG1heSBub3QgdXNlIHRoaXMgZmlsZSBleGNlcHQgaW4gY29tcGxpYW5jZSB3aXRoIHRoZSBMaWNlbnNlLlxuICogWW91IG1heSBvYnRhaW4gYSBjb3B5IG9mIHRoZSBMaWNlbnNlIGF0XG4gKlxuICogaHR0cDovL3d3dy5hcGFjaGUub3JnL2xpY2Vuc2VzL0xJQ0VOU0UtMi4wXG4gKlxuICogVW5sZXNzIHJlcXVpcmVkIGJ5IGFwcGxpY2FibGUgbGF3IG9yIGFncmVlZCB0byBpbiB3cml0aW5nLCBzb2Z0d2FyZVxuICogZGlzdHJpYnV0ZWQgdW5kZXIgdGhlIExpY2Vuc2UgaXMgZGlzdHJpYnV0ZWQgb24gYW4gXCJBUyBJU1wiIEJBU0lTLFxuICogV0lUSE9VVCBXQVJSQU5USUVTIE9SIENPTkRJVElPTlMgT0YgQU5ZIEtJTkQsIGVpdGhlciBleHByZXNzIG9yIGltcGxpZWQuXG4gKiBTZWUgdGhlIExpY2Vuc2UgZm9yIHRoZSBzcGVjaWZpYyBsYW5ndWFnZSBnb3Zlcm5pbmcgcGVybWlzc2lvbnMgYW5kXG4gKiBsaW1pdGF0aW9ucyB1bmRlciB0aGUgTGljZW5zZS5cbiAqID09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09PT09XG4gKi9cblxuaW1wb3J0IHtLZXJuZWxDb25maWcsIEtlcm5lbEZ1bmMsIFNlYXJjaFNvcnRlZCwgU2VhcmNoU29ydGVkQXR0cnMsIFNlYXJjaFNvcnRlZElucHV0cywgVGVuc29ySW5mb30gZnJvbSAnQHRlbnNvcmZsb3cvdGZqcy1jb3JlJztcblxuaW1wb3J0IHtCYWNrZW5kV2FzbX0gZnJvbSAnLi4vYmFja2VuZF93YXNtJztcbmltcG9ydCB7Q3BwRFR5cGV9IGZyb20gJy4vdHlwZXMnO1xuXG5sZXQgd2FzbVNlYXJjaFNvcnRlZDogKFxuICAgIHNvcnRlZFNlcXVlbmNlSWQ6IG51bWJlciwgdmFsdWVzSWQ6IG51bWJlciwgYmF0Y2hTaXplOiBudW1iZXIsXG4gICAgc2VxdWVuY2VTaXplOiBudW1iZXIsIHZhbHVlc1NpemU6IG51bWJlciwgZHR5cGU6IG51bWJlcixcbiAgICBpc1NpZGVMZWZ0OiBib29sZWFuLCBvdXRJZDogbnVtYmVyKSA9PiB2b2lkO1xuXG5mdW5jdGlvbiBzZXR1cChiYWNrZW5kOiBCYWNrZW5kV2FzbSkge1xuICB3YXNtU2VhcmNoU29ydGVkID0gYmFja2VuZC53YXNtLmN3cmFwKFNlYXJjaFNvcnRlZCwgbnVsbCAvKiB2b2lkICovLCBbXG4gICAgJ251bWJlcicsICAvLyBzb3J0ZWRTZXF1ZW5jZUlkXG4gICAgJ251bWJlcicsICAvLyB2YWx1ZXNJZFxuICAgICdudW1iZXInLCAgLy8gYmF0Y2hTaXplXG4gICAgJ251bWJlcicsICAvLyBzZXF1ZW5jZVNpemVcbiAgICAnbnVtYmVyJywgIC8vIHZhbHVlc1NpemVcbiAgICAnbnVtYmVyJywgIC8vIGR0eXBlXG4gICAgJ2Jvb2wnLCAgICAvLyBpc1NpZGVMZWZ0XG4gICAgJ251bWJlcicsICAvLyBvdXRJZFxuICBdKTtcbn1cblxuZnVuY3Rpb24gc2VhcmNoU29ydGVkKGFyZ3M6IHtcbiAgaW5wdXRzOiBTZWFyY2hTb3J0ZWRJbnB1dHMsXG4gIGJhY2tlbmQ6IEJhY2tlbmRXYXNtLFxuICBhdHRyczogU2VhcmNoU29ydGVkQXR0cnMsXG59KSB7XG4gIGNvbnN0IHtpbnB1dHMsIGJhY2tlbmQsIGF0dHJzfSA9IGFyZ3M7XG4gIGNvbnN0IHtzb3J0ZWRTZXF1ZW5jZSwgdmFsdWVzfSA9IGlucHV0cztcbiAgY29uc3Qge3NpZGV9ID0gYXR0cnM7XG5cbiAgaWYgKHNvcnRlZFNlcXVlbmNlLmR0eXBlICE9PSB2YWx1ZXMuZHR5cGUpIHtcbiAgICB0aHJvdyBuZXcgRXJyb3IoXG4gICAgICAgIGBTZWFyY2hTb3J0ZWQgZXJyb3I6IHNvcnRlZF9zZXF1ZW5jZSBtdXN0IGhhdmUgdGhlIHNhbWUgZHR5cGUgYXMgdmFsdWVzLiBHb3QgJHtcbiAgICAgICAgICAgIHNvcnRlZFNlcXVlbmNlLmR0eXBlfSBhbmQgJHt2YWx1ZXMuZHR5cGV9YCk7XG4gIH1cblxuICBjb25zdCBvdXQgPSBiYWNrZW5kLm1ha2VPdXRwdXQodmFsdWVzLnNoYXBlLCAnaW50MzInKTtcblxuICBmdW5jdGlvbiB0ZW5zb3JJZCh4OiBUZW5zb3JJbmZvKSB7XG4gICAgcmV0dXJuIGJhY2tlbmQuZGF0YUlkTWFwLmdldCh4LmRhdGFJZCkuaWQhO1xuICB9XG4gIHdhc21TZWFyY2hTb3J0ZWQoXG4gICAgICB0ZW5zb3JJZChzb3J0ZWRTZXF1ZW5jZSksXG4gICAgICB0ZW5zb3JJZCh2YWx1ZXMpLFxuICAgICAgLypiYXRjaFNpemU9Ki9zb3J0ZWRTZXF1ZW5jZS5zaGFwZVswXSxcbiAgICAgIC8qc2VxdWVuY2VTaXplPSovc29ydGVkU2VxdWVuY2Uuc2hhcGVbMV0sXG4gICAgICAvKnZhbHVlc1NpemU9Ki92YWx1ZXMuc2hhcGVbMV0sXG4gICAgICAvKmR0eXBlPSovQ3BwRFR5cGVbc29ydGVkU2VxdWVuY2UuZHR5cGVdLFxuICAgICAgLyppc1NpZGVMZWZ0PSovc2lkZSA9PT0gJ2xlZnQnLFxuICAgICAgdGVuc29ySWQob3V0KSxcbiAgKTtcblxuICByZXR1cm4gb3V0O1xufVxuXG5leHBvcnQgY29uc3Qgc2VhcmNoU29ydGVkQ29uZmlnOiBLZXJuZWxDb25maWcgPSB7XG4gIGtlcm5lbE5hbWU6IFNlYXJjaFNvcnRlZCxcbiAgYmFja2VuZE5hbWU6ICd3YXNtJyxcbiAgc2V0dXBGdW5jOiBzZXR1cCxcbiAga2VybmVsRnVuYzogc2VhcmNoU29ydGVkIGFzIHVua25vd24gYXMgS2VybmVsRnVuY1xufTtcbiJdfQ==