| 1 | {"version":3,"file":"split_util.js","sourceRoot":"","sources":["../../src/ops/split_util.ts"],"names":[],"mappings":"AAkBA,OAAO,EAAC,MAAM,EAAC,MAAM,SAAS,CAAC;AAE/B;;;;GAIG;AACH,MAAM,UAAU,gBAAgB,CAC5B,CAAoB,EAAE,eAAgC,EACtD,IAAI,GAAG,CAAC;IACV,IAAI,UAAU,GAAG,EAAE,CAAC;IACpB,IAAI,OAAO,CAAC,eAAe,CAAC,KAAK,QAAQ,EAAE;QACzC,MAAM,CACF,CAAC,CAAC,KAAK,CAAC,IAAI,CAAC,GAAG,eAAe,KAAK,CAAC,EACrC,GAAG,EAAE,CAAC,+CAA+C,CAAC,CAAC;QAC3D,UAAU;YACN,IAAI,KAAK,CAAC,eAAe,CAAC,CAAC,IAAI,CAAC,CAAC,CAAC,KAAK,CAAC,IAAI,CAAC,GAAG,eAAe,CAAC,CAAC;KACtE;SAAM;QACL,MAAM,SAAS,GAAG,eAAe,CAAC,MAAM,CAAC,CAAC,KAAK,EAAE,KAAK,EAAE,EAAE;YACxD,IAAI,KAAK,KAAK,CAAC,CAAC,EAAE;gBAChB,KAAK,IAAI,CAAC,CAAC;aACZ;YACD,OAAO,KAAK,CAAC;QACf,CAAC,EAAE,CAAC,CAAC,CAAC;QACN,MAAM,CACF,SAAS,IAAI,CAAC,EACd,GAAG,EAAE,CAAC,yDAAyD,CAAC,CAAC;QACrE,MAAM,QAAQ,GAAG,eAAe,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,CAAC;QAC7C,qEAAqE;QACrE,2CAA2C;QAC3C,IAAI,QAAQ,KAAK,CAAC,CAAC,EAAE;YACnB,MAAM,KAAK,GAAG,eAAe,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC,CAAC;YAClE,eAAe,CAAC,QAAQ,CAAC,GAAG,CAAC,CAAC,KAAK,CAAC,IAAI,CAAC,GAAG,KAAK,CAAC;SACnD;QACD,MAAM,CACF,CAAC,CAAC,KAAK,CAAC,IAAI,CAAC,KAAK,eAAe,CAAC,MAAM,CAAC,CAAC,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC,CAAC,GAAG,CAAC,CAAC,EACzD,GAAG,EAAE,CAAC,6DAA6D,CAAC,CAAC;QACzE,UAAU,GAAG,eAAe,CAAC;KAC9B;IAED,OAAO,UAAU,CAAC;AACpB,CAAC","sourcesContent":["/**\n * @license\n * Copyright 2020 Google LLC. All Rights Reserved.\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n * =============================================================================\n */\nimport {TensorInfo} from '../kernel_registry';\nimport {Tensor} from '../tensor';\nimport {assert} from '../util';\n\n/**\n * Prepare the split size array. When the input is a number, the axis is evenly\n * divided among the split size. When the input contains the negative value, the\n * rest of the axis is allocated toward that.\n */\nexport function prepareSplitSize(\n x: Tensor|TensorInfo, numOrSizeSplits: number[]|number,\n axis = 0): number[] {\n let splitSizes = [];\n if (typeof (numOrSizeSplits) === 'number') {\n assert(\n x.shape[axis] % numOrSizeSplits === 0,\n () => 'Number of splits must evenly divide the axis.');\n splitSizes =\n new Array(numOrSizeSplits).fill(x.shape[axis] / numOrSizeSplits);\n } else {\n const numOfNegs = numOrSizeSplits.reduce((count, value) => {\n if (value === -1) {\n count += 1;\n }\n return count;\n }, 0);\n assert(\n numOfNegs <= 1,\n () => 'There should be only one negative value in split array.');\n const negIndex = numOrSizeSplits.indexOf(-1);\n // Allow the number of split array to be -1, which indicates the rest\n // of dimension is allocated to that split.\n if (negIndex !== -1) {\n const total = numOrSizeSplits.reduce((a, b) => b > 0 ? a + b : a);\n numOrSizeSplits[negIndex] = x.shape[axis] - total;\n }\n assert(\n x.shape[axis] === numOrSizeSplits.reduce((a, b) => a + b),\n () => 'The sum of sizes must match the size of the axis dimension.');\n splitSizes = numOrSizeSplits;\n }\n\n return splitSizes;\n}\n"]} |