wehub-resource-sync
367 lines
13 KiB
JavaScript
367 lines
13 KiB
JavaScript
|
|
import * as fs from 'fs/promises';
|
|
import * as path from 'path';
|
|
import * as protobuf from '../source/protobuf.js';
|
|
import * as url from 'url';
|
|
import { tensorflow } from '../source/tf-proto.js';
|
|
|
|
const decoder = new TextDecoder('utf-8');
|
|
|
|
const has = (obj, key) => Object.prototype.hasOwnProperty.call(obj, key);
|
|
|
|
const dataTypes = new Map([
|
|
[tensorflow.DataType.DT_HALF, 'float16'],
|
|
[tensorflow.DataType.DT_FLOAT, 'float32'],
|
|
[tensorflow.DataType.DT_DOUBLE, 'float64'],
|
|
[tensorflow.DataType.DT_INT32, 'int32'],
|
|
[tensorflow.DataType.DT_UINT8, 'uint8'],
|
|
[tensorflow.DataType.DT_UINT16, 'uint16'],
|
|
[tensorflow.DataType.DT_UINT32, 'uint32'],
|
|
[tensorflow.DataType.DT_UINT64, 'uint64'],
|
|
[tensorflow.DataType.DT_INT16, 'int16'],
|
|
[tensorflow.DataType.DT_INT8, 'int8'],
|
|
[tensorflow.DataType.DT_STRING, 'string'],
|
|
[tensorflow.DataType.DT_COMPLEX64, 'complex64'],
|
|
[tensorflow.DataType.DT_COMPLEX128, 'complex128'],
|
|
[tensorflow.DataType.DT_INT64, 'int64'],
|
|
[tensorflow.DataType.DT_BOOL, 'bool'],
|
|
[tensorflow.DataType.DT_QINT8, 'qint8'],
|
|
[tensorflow.DataType.DT_QUINT8, 'quint8'],
|
|
[tensorflow.DataType.DT_QINT16, 'qint16'],
|
|
[tensorflow.DataType.DT_QUINT16, 'quint16'],
|
|
[tensorflow.DataType.DT_QINT32, 'qint32'],
|
|
[tensorflow.DataType.DT_BFLOAT16, 'bfloat16'],
|
|
[tensorflow.DataType.DT_RESOURCE, 'resource'],
|
|
[tensorflow.DataType.DT_VARIANT, 'variant'],
|
|
[tensorflow.DataType.DT_HALF_REF, 'float16_ref'],
|
|
[tensorflow.DataType.DT_FLOAT_REF, 'float32_ref'],
|
|
[tensorflow.DataType.DT_DOUBLE_REF, 'float64_ref'],
|
|
[tensorflow.DataType.DT_INT32_REF, 'int32_ref'],
|
|
[tensorflow.DataType.DT_UINT32_REF, 'uint32_ref'],
|
|
[tensorflow.DataType.DT_UINT8_REF, 'uint8_ref'],
|
|
[tensorflow.DataType.DT_UINT16_REF, 'uint16_ref'],
|
|
[tensorflow.DataType.DT_INT16_REF, 'int16_ref'],
|
|
[tensorflow.DataType.DT_INT8_REF, 'int8_ref'],
|
|
[tensorflow.DataType.DT_STRING_REF, 'string_ref'],
|
|
[tensorflow.DataType.DT_COMPLEX64_REF, 'complex64_ref'],
|
|
[tensorflow.DataType.DT_COMPLEX128_REF, 'complex128_ref'],
|
|
[tensorflow.DataType.DT_INT64_REF, 'int64_ref'],
|
|
[tensorflow.DataType.DT_UINT64_REF, 'uint64_ref'],
|
|
[tensorflow.DataType.DT_BOOL_REF, 'bool_ref'],
|
|
[tensorflow.DataType.DT_QINT8_REF, 'qint8_ref'],
|
|
[tensorflow.DataType.DT_QUINT8_REF, 'quint8_ref'],
|
|
[tensorflow.DataType.DT_QINT16_REF, 'qint16_ref'],
|
|
[tensorflow.DataType.DT_QUINT16_REF, 'quint16_ref'],
|
|
[tensorflow.DataType.DT_QINT32_REF, 'qint32_ref'],
|
|
[tensorflow.DataType.DT_BFLOAT16_REF, 'bfloat16_ref'],
|
|
[tensorflow.DataType.DT_RESOURCE_REF, 'resource_ref'],
|
|
[tensorflow.DataType.DT_VARIANT_REF, 'variant_ref']
|
|
]);
|
|
|
|
const attributeTypes = new Map([
|
|
['type', 'type'], ['list(type)', 'type[]'],
|
|
['bool', 'boolean'],
|
|
['int', 'int64'], ['list(int)', 'int64[]'],
|
|
['float', 'float32'], ['list(float)', 'float32[]'],
|
|
['string', 'string'], ['list(string)', 'string[]'],
|
|
['shape', 'shape'], ['list(shape)', 'shape[]'],
|
|
['tensor', 'tensor'],
|
|
['func', 'function'], ['list(func)', 'function[]']
|
|
]);
|
|
|
|
const categories = new Map([
|
|
['Assign', 'Control'],
|
|
['AvgPool', 'Pool'],
|
|
['BatchNormWithGlobalNormalization', 'Normalization'],
|
|
['BiasAdd', 'Layer'],
|
|
['Concat', 'Tensor'],
|
|
['ConcatV2', 'Tensor'],
|
|
['Const', 'Constant'],
|
|
['Conv2D', 'Layer'],
|
|
['DepthwiseConv2dNative', 'Layer'],
|
|
['Dequantize', 'Quantization'],
|
|
['Elu', 'Activation'],
|
|
['FusedBatchNorm', 'Normalization'],
|
|
['FusedBatchNormV2', 'Normalization'],
|
|
['FusedBatchNormV3', 'Normalization'],
|
|
['Gather', 'Transform'],
|
|
['Identity', 'Control'],
|
|
['LeakyRelu', 'Activation'],
|
|
['LRN', 'Normalization'],
|
|
['LSTMBlockCell', 'Layer'],
|
|
['MaxPool', 'Pool'],
|
|
['MaxPoolV2', 'Pool'],
|
|
['MaxPoolWithArgmax', 'Pool'],
|
|
['Pad', 'Tensor'],
|
|
['QuantizeAndDequantize', 'Quantization'],
|
|
['QuantizeAndDequantizeV2', 'Quantization'],
|
|
['QuantizeAndDequantizeV3', 'Quantization'],
|
|
['QuantizeAndDequantizeV4', 'Quantization'],
|
|
['QuantizeAndDequantizeV4Grad', 'Quantization'],
|
|
['QuantizeDownAndShrinkRange', 'Quantization'],
|
|
['QuantizeV2', 'Quantization'],
|
|
['Relu', 'Activation'],
|
|
['Relu6', 'Activation'],
|
|
['Reshape', 'Shape'],
|
|
['Sigmoid', 'Activation'],
|
|
['Slice', 'Tensor'],
|
|
['Softmax', 'Activation'],
|
|
['Split', 'Tensor'],
|
|
['Squeeze', 'Transform'],
|
|
['StridedSlice', 'Tensor'],
|
|
['swish_f32', 'Activation'],
|
|
['Transpose', 'Transform'],
|
|
['Variable', 'Control'],
|
|
['VariableV2', 'Control']
|
|
]);
|
|
|
|
const convertFloat = (value) => {
|
|
if (value === Infinity) {
|
|
return 'NaN';
|
|
}
|
|
if (value === -Infinity) {
|
|
return '-NaN';
|
|
}
|
|
return Math.fround(value);
|
|
};
|
|
|
|
const escapeString = (text) => {
|
|
let result = '';
|
|
for (const c of text) {
|
|
switch (c) {
|
|
case '\n': result += '\\n'; break;
|
|
case '\r': result += '\\r'; break;
|
|
case '\t': result += '\\t'; break;
|
|
case '"': result += '\\"'; break;
|
|
case "'": result += "\\'"; break;
|
|
case '\\': result += '\\\\'; break;
|
|
default: result += c; break;
|
|
}
|
|
}
|
|
return result;
|
|
};
|
|
|
|
const heredoc = (input) => {
|
|
const lines = input.split('\n');
|
|
const output = [];
|
|
for (let i = 0; i < lines.length; i++) {
|
|
const line = lines[i];
|
|
const colon = line.indexOf(':');
|
|
const marker = colon === -1 ? null : line.substring(colon + 1).replace(/^ +/, '');
|
|
if (marker === null || !marker.startsWith('<<')) {
|
|
output.push(line);
|
|
continue;
|
|
}
|
|
const terminator = marker.substring(2);
|
|
const content = [];
|
|
let trailing = '';
|
|
while (++i < lines.length) {
|
|
const inner = lines[i];
|
|
if (inner.startsWith(terminator)) {
|
|
trailing = inner.substring(terminator.length);
|
|
break;
|
|
}
|
|
content.push(inner);
|
|
}
|
|
output.push(`${line.substring(0, colon + 1)}"${escapeString(content.join('\n'))}"${trailing}`);
|
|
}
|
|
return output.join('\n');
|
|
};
|
|
|
|
const readOpList = async (file) => {
|
|
const content = await fs.readFile(file, 'utf-8');
|
|
const encoder = new TextEncoder();
|
|
const reader = new protobuf.TextReader(encoder.encode(content));
|
|
return tensorflow.OpList.decodeText(reader);
|
|
};
|
|
|
|
const readApiDefs = async (folder) => {
|
|
const dirs = await fs.readdir(folder);
|
|
const files = dirs.filter((name) => name.endsWith('.pbtxt'));
|
|
const encoder = new TextEncoder();
|
|
const defs = new Map();
|
|
for (const name of files.sort()) {
|
|
const file = path.join(folder, name);
|
|
// eslint-disable-next-line no-await-in-loop
|
|
const content = await fs.readFile(file, 'utf-8');
|
|
const text = heredoc(content);
|
|
const buffer = encoder.encode(text);
|
|
const reader = protobuf.TextReader.open(buffer);
|
|
const apiDefs = tensorflow.ApiDefs.decodeText(reader);
|
|
for (const op of apiDefs.op) {
|
|
defs.set(op.graph_op_name, op);
|
|
}
|
|
}
|
|
return defs;
|
|
};
|
|
|
|
const convertAttrList = (attrValue) => {
|
|
const result = [];
|
|
const list = attrValue.list;
|
|
for (const value of list.s) {
|
|
result.push(decoder.decode(value));
|
|
}
|
|
for (const value of list.i) {
|
|
result.push(typeof value === 'bigint' ? Number(value) : value);
|
|
}
|
|
for (const value of list.f) {
|
|
result.push(convertFloat(value));
|
|
}
|
|
for (const value of list.type) {
|
|
result.push({ type: 'type', value });
|
|
}
|
|
if (result.length === 0 && (list.b.length > 0 || list.shape.length > 0 ||
|
|
list.tensor.length > 0 || list.func.length > 0)) {
|
|
throw new Error('Unsupported list value.');
|
|
}
|
|
return result;
|
|
};
|
|
|
|
const convertAttrValue = (attrValue) => {
|
|
if (has(attrValue, 'list')) {
|
|
return convertAttrList(attrValue);
|
|
}
|
|
if (has(attrValue, 's')) {
|
|
return decoder.decode(attrValue.s);
|
|
}
|
|
if (has(attrValue, 'i')) {
|
|
return typeof attrValue.i === 'bigint' ? Number(attrValue.i) : attrValue.i;
|
|
}
|
|
if (has(attrValue, 'f')) {
|
|
return convertFloat(attrValue.f);
|
|
}
|
|
if (has(attrValue, 'b')) {
|
|
return attrValue.b;
|
|
}
|
|
if (has(attrValue, 'type')) {
|
|
return { type: 'type', value: attrValue.type };
|
|
}
|
|
if (has(attrValue, 'tensor')) {
|
|
return { type: 'tensor', value: '?' };
|
|
}
|
|
if (has(attrValue, 'shape')) {
|
|
return { type: 'shape', value: '?' };
|
|
}
|
|
throw new Error('Unsupported attribute value.');
|
|
};
|
|
|
|
const formatAttributeValue = (value) => {
|
|
if (value && typeof value === 'object' && value.type === 'type') {
|
|
if (!dataTypes.has(value.value)) {
|
|
throw new Error(`Unknown data type '${value.value}'.`);
|
|
}
|
|
return dataTypes.get(value.value);
|
|
}
|
|
if (typeof value === 'string') {
|
|
return value;
|
|
}
|
|
if (value === true || value === false) {
|
|
return value.toString();
|
|
}
|
|
throw new Error('Unsupported attribute value.');
|
|
};
|
|
|
|
const buildSchema = (operator, apiDef) => {
|
|
const schema = { name: operator.name };
|
|
if (categories.has(operator.name)) {
|
|
schema.category = categories.get(operator.name);
|
|
}
|
|
if (apiDef.summary) {
|
|
schema.summary = apiDef.summary;
|
|
}
|
|
if (apiDef.description) {
|
|
schema.description = apiDef.description;
|
|
}
|
|
const apiDefAttr = new Map(apiDef.attr.map((attr) => [attr.name, attr]));
|
|
for (const attr of operator.attr) {
|
|
if (!attributeTypes.has(attr.type)) {
|
|
throw new Error(`Unknown attribute type '${attr.type}'.`);
|
|
}
|
|
const json = { name: attr.name, type: attributeTypes.get(attr.type) };
|
|
const description = apiDefAttr.get(attr.name)?.description;
|
|
if (description) {
|
|
json.description = description;
|
|
}
|
|
if (attr.has_minimum) {
|
|
json.minimum = Number(attr.minimum);
|
|
}
|
|
if (attr.allowed_values !== null) {
|
|
const allowed = convertAttrValue(attr.allowed_values).map((v) => `\`${formatAttributeValue(v)}\``).join(', ');
|
|
const prefix = json.description ? `${json.description} ` : '';
|
|
json.description = `${prefix}Must be one of the following: ${allowed}.`;
|
|
}
|
|
if (attr.default_value !== null) {
|
|
json.default = convertAttrValue(attr.default_value);
|
|
}
|
|
schema.attributes = schema.attributes || [];
|
|
schema.attributes.push(json);
|
|
}
|
|
const apiDefIn = new Map(apiDef.in_arg.map((arg) => [arg.name, arg]));
|
|
for (const arg of operator.input_arg) {
|
|
const json = { name: arg.name };
|
|
const description = apiDefIn.get(arg.name)?.description;
|
|
if (description) {
|
|
json.description = description;
|
|
}
|
|
if (arg.number_attr) {
|
|
json.numberAttr = arg.number_attr;
|
|
}
|
|
if (arg.type) {
|
|
json.type = arg.type;
|
|
}
|
|
if (arg.type_attr) {
|
|
json.typeAttr = arg.type_attr;
|
|
}
|
|
if (arg.type_list_attr) {
|
|
json.typeListAttr = arg.type_list_attr;
|
|
}
|
|
if (arg.is_ref) {
|
|
json.isRef = true;
|
|
}
|
|
schema.inputs = schema.inputs || [];
|
|
schema.inputs.push(json);
|
|
}
|
|
const apiDefOut = new Map(apiDef.out_arg.map((arg) => [arg.name, arg]));
|
|
for (const arg of operator.output_arg) {
|
|
const json = { name: arg.name };
|
|
const description = apiDefOut.get(arg.name)?.description;
|
|
if (description) {
|
|
json.description = description;
|
|
}
|
|
if (arg.number_attr) {
|
|
json.numberAttr = arg.number_attr;
|
|
}
|
|
if (arg.type) {
|
|
json.type = arg.type;
|
|
} else if (arg.type_attr) {
|
|
json.typeAttr = arg.type_attr;
|
|
} else if (arg.type_list_attr) {
|
|
json.typeListAttr = arg.type_list_attr;
|
|
}
|
|
if (arg.is_ref) {
|
|
json.isRef = true;
|
|
}
|
|
schema.outputs = schema.outputs || [];
|
|
schema.outputs.push(json);
|
|
}
|
|
return schema;
|
|
};
|
|
|
|
const main = async () => {
|
|
const root = path.dirname(path.dirname(url.fileURLToPath(import.meta.url)));
|
|
const dir = path.join(root, 'third_party', 'source', 'tensorflow', 'tensorflow', 'core');
|
|
const defs = await readApiDefs(path.join(dir, 'api_def', 'base_api'));
|
|
const ops = await readOpList(path.join(dir, 'ops', 'ops.pbtxt'));
|
|
const schemas = [];
|
|
for (const op of ops.op) {
|
|
const apiDef = defs.get(op.name) || new tensorflow.ApiDef();
|
|
const schema = buildSchema(op, apiDef);
|
|
schemas.push(schema);
|
|
}
|
|
const file = path.join(root, 'source', 'tf-metadata.json');
|
|
const content = JSON.stringify(schemas, null, 2);
|
|
await fs.writeFile(file, content, 'utf-8');
|
|
};
|
|
|
|
await main();
|