Function getTensorStorages

torch/utils/model_dump/code.js:483–534

(data)

// TODO: Maybe track by dtype as well.
// TODO: Maybe distinguish between visible size and storage size.
function getTensorStorages(data) {
  if (data === null) {
    return new Map();
  }
  if (typeof(data) == "boolean") {
    return new Map();
  }
  if (typeof(data) == "number") {
    return new Map();
  }
  if (typeof(data) == "string") {
    return new Map();
  }
  if (typeof(data) != "object") {
    throw new Error("Not an object");
  }
  if (Array.isArray(data)) {
    let result = new Map();
    for (const item of data) {
      const tensors = getTensorStorages(item);
      for (const [key, storage] of tensors.entries()) {
        if (!result.has(key)) {
          result.set(key, storage);
        } else {
          const old_storage = result.get(key);
          assertStorageAreEqual(key, old_storage, storage);
        }
      }
    }
    return result;
  }
  if (data.__tuple_values__) {
    return getTensorStorages(data.__tuple_values__);
  }
  if (data.__is_dict__) {
    return getTensorStorages(data.values);
  }
  if (data.__module_type__) {
    return getTensorStorages(data.state);
  }
  if (data.__tensor_v2__) {
    const [storage, offset, size, stride, grad] = data.__tensor_v2__;
    const [dtype, key, device, numel] = storage;
    return new Map([[key, storage]]);
  }
  if (data.__qtensor__) {
    const [storage, offset, size, stride, quantizer, grad] = data.__qtensor__;
    const [dtype, key, device, numel] = storage;
    return new Map([[key, storage]]);
  }
  throw new Error("Can't handle data type.", data);
}

function getTensorMemoryByDevice(pickles) {
  let all_tensors = [];