@tensorflow/tfjs-core/src/tensor_util.ts

Hardware-accelerated JavaScript library for machine intelligence

/**
 * @license
 * Copyright 2018 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 {Tensor} from './tensor';
import {TensorContainer, TensorContainerArray} from './tensor_types';
import {upcastType} from './types';
import {assert} from './util';

export function makeTypesMatch<T extends Tensor>(a: T, b: T): [T, T] {
  if (a.dtype === b.dtype) {
    return [a, b];
  }
  const dtype = upcastType(a.dtype, b.dtype);
  return [a.cast(dtype), b.cast(dtype)];
}

export function assertTypesMatch(a: Tensor, b: Tensor): void {
  assert(
      a.dtype === b.dtype,
      () => `The dtypes of the first(${a.dtype}) and` +
          ` second(${b.dtype}) input must match`);
}

export function isTensorInList(tensor: Tensor, tensorList: Tensor[]): boolean {
  return tensorList.some(x => x.id === tensor.id);
}

/**
 * Extracts any `Tensor`s found within the provided object.
 *
 * @param container an object that may be a `Tensor` or may directly contain
 *   `Tensor`s, such as a `Tensor[]` or `{key: Tensor, ...}`. In general it
 *   is safe to pass any object here, except that `Promise`s are not
 *   supported.
 * @returns An array of `Tensors` found within the passed object. If the
 *   argument is simply a `Tensor', a list containing that `Tensor` is
 *   returned. If the object is not a `Tensor` or does not
 *   contain `Tensors`, an empty list is returned.
 */
export function getTensorsInContainer(result: TensorContainer): Tensor[] {
  const list: Tensor[] = [];
  const seen = new Set<{}|void>();
  walkTensorContainer(result, list, seen);
  return list;
}

function walkTensorContainer(
    container: TensorContainer, list: Tensor[], seen: Set<{}|void>): void {
  if (container == null) {
    return;
  }
  if (container instanceof Tensor) {
    list.push(container);
    return;
  }
  if (!isIterable(container)) {
    return;
  }
  // Iteration over keys works also for arrays.
  const iterable = container as TensorContainerArray;
  for (const k in iterable) {
    const val = iterable[k];
    if (!seen.has(val)) {
      seen.add(val);
      walkTensorContainer(val, list, seen);
    }
  }
}

// tslint:disable-next-line:no-any
function isIterable(obj: any): boolean {
  return Array.isArray(obj) || typeof obj === 'object';
}