Function inferShape

tfjs-core/src/tensor_util_env.ts:25–55

(
    val: TensorLike|WebGLData|WebGPUData, dtype?: DataType)

import {bytesPerElement} from './util_base';

export function inferShape(
    val: TensorLike|WebGLData|WebGPUData, dtype?: DataType): number[] {
  let firstElem: typeof val = val;

  if (isTypedArray(val)) {
    return dtype === 'string' ? [] : [val.length];
  }

  if (isWebGLData(val)) {
    const usedChannels = val.channels || 'RGBA';
    return [val.height, val.width * usedChannels.length];
  } else if (isWebGPUData(val)) {
    return [val.buffer.size / (dtype == null ? 4 : bytesPerElement(dtype))];
  }
  if (!Array.isArray(val)) {
    return [];  // Scalar.
  }
  const shape: number[] = [];

  while (Array.isArray(firstElem) ||
         isTypedArray(firstElem) && dtype !== 'string') {
    shape.push(firstElem.length);
    firstElem = firstElem[0];
  }
  if (Array.isArray(val) &&
      env().getBool('TENSORLIKE_CHECK_SHAPE_CONSISTENCY')) {
    deepAssertShapeConsistency(val, shape, []);
  }

  return shape;
}

function deepAssertShapeConsistency(
    val: TensorLike, shape: number[], indices: number[]) {