| 1 | {"version":3,"file":"where.js","sourceRoot":"","sources":["../../src/ops/where.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;GAeG;AAEH,OAAO,EAAC,MAAM,EAAC,MAAM,WAAW,CAAC;AACjC,OAAO,EAAC,MAAM,EAAe,MAAM,iBAAiB,CAAC;AAGrD,OAAO,EAAC,eAAe,EAAC,MAAM,oBAAoB,CAAC;AAGnD,OAAO,EAAC,WAAW,EAAC,MAAM,gBAAgB,CAAC;AAC3C,OAAO,EAAC,0BAA0B,EAAC,MAAM,kBAAkB,CAAC;AAC5D,OAAO,EAAC,EAAE,EAAC,MAAM,aAAa,CAAC;AAE/B;;;;;;;;;;;;;;;;;;;;;;GAsBG;AACH,SAAS,MAAM,CACX,SAA4B,EAAE,CAAe,EAAE,CAAe;IAChE,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,GAAG,EAAE,OAAO,CAAC,CAAC;IAC5C,MAAM,EAAE,GAAG,eAAe,CAAC,CAAC,EAAE,GAAG,EAAE,OAAO,CAAC,CAAC;IAC5C,MAAM,UAAU,GAAG,eAAe,CAAC,SAAS,EAAE,WAAW,EAAE,OAAO,EAAE,MAAM,CAAC,CAAC;IAC5E,uEAAuE;IACvE,uBAAuB;IACvB,2DAA2D;IAC3D,MAAM,cAAc,GAAG,0BAA0B,CAC7C,0BAA0B,CAAC,UAAU,CAAC,KAAK,EAAE,EAAE,CAAC,KAAK,CAAC,EAAE,EAAE,CAAC,KAAK,CAAC,CAAC;IACtE,MAAM,qBAAqB,GAAG,WAAW,CAAC,UAAU,EAAE,cAAc,CAAC,CAAC;IACtE,MAAM,aAAa,GAAG,WAAW,CAAC,EAAE,EAAE,cAAc,CAAC,CAAC;IACtD,MAAM,aAAa,GAAG,WAAW,CAAC,EAAE,EAAE,cAAc,CAAC,CAAC;IAEtD,MAAM,MAAM,GAAiB;QAC3B,SAAS,EAAE,qBAAqB;QAChC,CAAC,EAAE,aAAa;QAChB,CAAC,EAAE,aAAa;KACjB,CAAC;IACF,OAAO,MAAM,CAAC,SAAS,CAAC,MAAM,EAAE,MAA8B,CAAC,CAAC;AAClE,CAAC;AAED,MAAM,CAAC,MAAM,KAAK,GAAG,EAAE,CAAC,EAAC,MAAM,EAAC,CAAC,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 */\n\nimport {ENGINE} from '../engine';\nimport {Select, SelectInputs} from '../kernel_names';\nimport {Tensor} from '../tensor';\nimport {NamedTensorMap} from '../tensor_types';\nimport {convertToTensor} from '../tensor_util_env';\nimport {TensorLike} from '../types';\n\nimport {broadcastTo} from './broadcast_to';\nimport {assertAndGetBroadcastShape} from './broadcast_util';\nimport {op} from './operation';\n\n/**\n * Returns the elements, either `a` or `b` depending on the `condition`.\n *\n * If the condition is true, select from `a`, otherwise select from `b`.\n *\n * ```js\n * const cond = tf.tensor1d([false, false, true], 'bool');\n * const a = tf.tensor1d([1 , 2, 3]);\n * const b = tf.tensor1d([-1, -2, -3]);\n *\n * a.where(cond, b).print();\n * ```\n *\n * @param condition The input condition. Must be of dtype bool.\n * @param a If `condition` is rank 1, `a` may have a higher rank but\n * its first dimension must match the size of `condition`.\n * @param b A tensor with the same dtype as `a` and with shape that is\n * compatible with `a`.\n * @return A tensor with same dtype as `a` and `b`, and shape that is\n * broadcastable from `a` and `b`.\n *\n * @doc {heading: 'Operations', subheading: 'Logical'}\n */\nfunction where_<T extends Tensor>(\n condition: Tensor|TensorLike, a: T|TensorLike, b: T|TensorLike): T {\n const $a = convertToTensor(a, 'a', 'where');\n const $b = convertToTensor(b, 'b', 'where');\n const $condition = convertToTensor(condition, 'condition', 'where', 'bool');\n // TODO: move this logic to forward function when the broadcastTo op is\n // implemented in WASM.\n // Find the broadcastable shape for $condition, $a, and $b.\n const broadcastShape = assertAndGetBroadcastShape(\n assertAndGetBroadcastShape($condition.shape, $a.shape), $b.shape);\n const $broadcastedCondition = broadcastTo($condition, broadcastShape);\n const $broadcastedA = broadcastTo($a, broadcastShape);\n const $broadcastedB = broadcastTo($b, broadcastShape);\n\n const inputs: SelectInputs = {\n condition: $broadcastedCondition,\n t: $broadcastedA,\n e: $broadcastedB\n };\n return ENGINE.runKernel(Select, inputs as {} as NamedTensorMap);\n}\n\nexport const where = op({where_});\n"]} |