loadIntoBytes static method

void loadIntoBytes(Module module, Uint8List bytes)

Load a checkpoint from a byte buffer into an existing module. The module must have the same number of parameters with matching shapes in the same order.

Implementation

static void loadIntoBytes(Module module, Uint8List bytes) {
  if (bytes.length < 12) {
    throw ArgumentError(
      'Checkpoint: buffer too small (${bytes.length} bytes)',
    );
  }
  for (int i = 0; i < 4; i++) {
    if (bytes[i] != magic[i]) {
      throw ArgumentError(
        'Checkpoint: bad magic bytes; not a DPTC checkpoint',
      );
    }
  }
  final view = ByteData.sublistView(bytes);
  final ver = view.getUint32(4, Endian.little);
  if (ver != version) {
    throw ArgumentError(
      'Checkpoint: unsupported version $ver (expected $version)',
    );
  }
  final headerLen = view.getUint32(8, Endian.little);
  final headerEnd = 12 + headerLen;
  if (bytes.length < headerEnd) {
    throw ArgumentError(
      'Checkpoint: truncated header (need $headerEnd bytes, have '
      '${bytes.length})',
    );
  }
  final header =
      jsonDecode(utf8.decode(bytes.sublist(12, headerEnd)))
          as Map<String, dynamic>;
  final paramSpecs = (header['params'] as List).cast<Map<String, dynamic>>();
  final params = module.parameters();
  if (paramSpecs.length != params.length) {
    throw ArgumentError(
      'Checkpoint: parameter count mismatch — checkpoint has '
      '${paramSpecs.length}, module has ${params.length}',
    );
  }

  // Verify all shapes match before writing anything (avoid partial
  // loads).
  final paramDtypes = <DType>[];
  for (int i = 0; i < params.length; i++) {
    final want = (paramSpecs[i]['shape'] as List).cast<int>();
    final have = params[i].shape;
    if (want.length != have.length ||
        !List<bool>.generate(
          want.length,
          (k) => want[k] == have[k],
        ).every((b) => b)) {
      throw ArgumentError(
        'Checkpoint: shape mismatch for parameter #$i — checkpoint '
        '$want vs module $have',
      );
    }
    final dtypeStr = paramSpecs[i]['dtype'] as String?;
    switch (dtypeStr) {
      case null:
      case 'F32':
        paramDtypes.add(DType.fp32);
      case 'F16':
        paramDtypes.add(DType.fp16);
      default:
        throw ArgumentError(
          'Checkpoint: unsupported dtype "$dtypeStr" for parameter '
          '#$i (expected F32 or F16)',
        );
    }
  }
  final expectedDataLen = () {
    var n = 0;
    for (int i = 0; i < params.length; i++) {
      n += params[i].length * paramDtypes[i].itemBytes;
    }
    return n;
  }();
  if (bytes.length - headerEnd != expectedDataLen) {
    throw ArgumentError(
      'Checkpoint: data blob length ${bytes.length - headerEnd} '
      'does not match expected $expectedDataLen bytes',
    );
  }

  // Copy in.
  var off = headerEnd;
  for (int i = 0; i < params.length; i++) {
    final p = params[i];
    final n = p.length;
    final f32 = Float32List(n);
    if (paramDtypes[i] == DType.fp16) {
      for (int j = 0; j < n; j++) {
        f32[j] = fp16BitsToFp32(view.getUint16(off, Endian.little));
        off += 2;
      }
    } else {
      for (int j = 0; j < n; j++) {
        f32[j] = view.getFloat32(off, Endian.little);
        off += 4;
      }
    }
    final source = Tensor.fromList(p.shape, f32, device: p.device);
    p.assign(source);
  }
}