gguf.go 13.4 KB
Newer Older
Michael Yang's avatar
Michael Yang committed
1
package ggml
Bruce MacDonald's avatar
Bruce MacDonald committed
2
3
4

import (
	"bytes"
Michael Yang's avatar
Michael Yang committed
5
	"cmp"
Bruce MacDonald's avatar
Bruce MacDonald committed
6
	"encoding/binary"
7
	"encoding/json"
Bruce MacDonald's avatar
Bruce MacDonald committed
8
9
	"fmt"
	"io"
Michael Yang's avatar
Michael Yang committed
10
	"log/slog"
Michael Yang's avatar
Michael Yang committed
11
	"maps"
Michael Yang's avatar
Michael Yang committed
12
	"slices"
13
	"strings"
Bruce MacDonald's avatar
Bruce MacDonald committed
14
15
)

16
type containerGGUF struct {
17
	ByteOrder binary.ByteOrder
Michael Yang's avatar
ggufv3  
Michael Yang committed
18

Bruce MacDonald's avatar
Bruce MacDonald committed
19
20
21
22
23
24
25
26
27
28
29
	Version uint32

	V1 struct {
		NumTensor uint32
		NumKV     uint32
	}

	V2 struct {
		NumTensor uint64
		NumKV     uint64
	}
30
31
32
33
34

	V3 struct {
		NumTensor uint64
		NumKV     uint64
	}
35
36
37
38
39
40

	maxArraySize int
}

func (c *containerGGUF) canCollectArray(size int) bool {
	return c.maxArraySize < 0 || size <= c.maxArraySize
Bruce MacDonald's avatar
Bruce MacDonald committed
41
42
}

43
func (c *containerGGUF) Name() string {
Bruce MacDonald's avatar
Bruce MacDonald committed
44
45
46
	return "gguf"
}

47
func (c *containerGGUF) Decode(rs io.ReadSeeker) (model, error) {
48
	if err := binary.Read(rs, c.ByteOrder, &c.Version); err != nil {
49
50
		return nil, err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
51

52
	var err error
Bruce MacDonald's avatar
Bruce MacDonald committed
53
54
	switch c.Version {
	case 1:
55
56
57
		err = binary.Read(rs, c.ByteOrder, &c.V1)
	case 2:
		err = binary.Read(rs, c.ByteOrder, &c.V2)
Bruce MacDonald's avatar
Bruce MacDonald committed
58
	default:
59
60
61
62
		err = binary.Read(rs, c.ByteOrder, &c.V3)
	}
	if err != nil {
		return nil, err
Bruce MacDonald's avatar
Bruce MacDonald committed
63
64
	}

65
	model := newGGUF(c)
Michael Yang's avatar
Michael Yang committed
66
	if err := model.Decode(rs); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
67
68
69
70
71
72
		return nil, err
	}

	return model, nil
}

73
const (
74
75
76
77
78
79
80
81
82
83
84
85
86
	ggufTypeUint8 uint32 = iota
	ggufTypeInt8
	ggufTypeUint16
	ggufTypeInt16
	ggufTypeUint32
	ggufTypeInt32
	ggufTypeFloat32
	ggufTypeBool
	ggufTypeString
	ggufTypeArray
	ggufTypeUint64
	ggufTypeInt64
	ggufTypeFloat64
87
)
Bruce MacDonald's avatar
Bruce MacDonald committed
88

89
90
type gguf struct {
	*containerGGUF
91

Michael Yang's avatar
Michael Yang committed
92
93
	kv      KV
	tensors []*Tensor
94

Michael Yang's avatar
Michael Yang committed
95
	parameters   uint64
96
	tensorOffset uint64
97
98

	scratch [16 << 10]byte
Bruce MacDonald's avatar
Bruce MacDonald committed
99
100
}

101
102
103
func newGGUF(container *containerGGUF) *gguf {
	return &gguf{
		containerGGUF: container,
Michael Yang's avatar
Michael Yang committed
104
		kv:            make(KV),
Bruce MacDonald's avatar
Bruce MacDonald committed
105
106
107
	}
}

Michael Yang's avatar
Michael Yang committed
108
109
110
111
func (llm *gguf) KV() KV {
	return llm.kv
}

Michael Yang's avatar
Michael Yang committed
112
113
114
func (llm *gguf) Tensors() Tensors {
	return Tensors{
		items:  llm.tensors,
115
116
		Offset: llm.tensorOffset,
	}
Michael Yang's avatar
Michael Yang committed
117
118
}

119
120
121
func (llm *gguf) numTensor() uint64 {
	switch llm.Version {
	case 1:
122
		return uint64(llm.V1.NumTensor)
123
124
125
126
	case 2:
		return llm.V2.NumTensor
	default:
		return llm.V3.NumTensor
127
128
129
	}
}

130
131
132
func (llm *gguf) numKV() uint64 {
	switch llm.Version {
	case 1:
Bruce MacDonald's avatar
Bruce MacDonald committed
133
		return uint64(llm.V1.NumKV)
134
135
136
137
	case 2:
		return llm.V2.NumKV
	default:
		return llm.V3.NumKV
Bruce MacDonald's avatar
Bruce MacDonald committed
138
139
140
	}
}

141
142
143
144
145
func (llm *gguf) Decode(rs io.ReadSeeker) error {
	// decode key-values
	for i := 0; uint64(i) < llm.numKV(); i++ {
		k, err := readGGUFString(llm, rs)
		if err != nil {
146
147
148
			return err
		}

149
150
		t, err := readGGUF[uint32](llm, rs)
		if err != nil {
151
152
153
			return err
		}

154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
		var v any
		switch t {
		case ggufTypeUint8:
			v, err = readGGUF[uint8](llm, rs)
		case ggufTypeInt8:
			v, err = readGGUF[int8](llm, rs)
		case ggufTypeUint16:
			v, err = readGGUF[uint16](llm, rs)
		case ggufTypeInt16:
			v, err = readGGUF[int16](llm, rs)
		case ggufTypeUint32:
			v, err = readGGUF[uint32](llm, rs)
		case ggufTypeInt32:
			v, err = readGGUF[int32](llm, rs)
		case ggufTypeUint64:
			v, err = readGGUF[uint64](llm, rs)
		case ggufTypeInt64:
			v, err = readGGUF[int64](llm, rs)
		case ggufTypeFloat32:
			v, err = readGGUF[float32](llm, rs)
		case ggufTypeFloat64:
			v, err = readGGUF[float64](llm, rs)
		case ggufTypeBool:
			v, err = readGGUF[bool](llm, rs)
		case ggufTypeString:
			v, err = readGGUFString(llm, rs)
		case ggufTypeArray:
			v, err = readGGUFArray(llm, rs)
		default:
			return fmt.Errorf("invalid type: %d", t)
184
185
		}

186
		if err != nil {
187
188
189
			return err
		}

Michael Yang's avatar
Michael Yang committed
190
		llm.kv[k] = v
191
192
	}

193
	// decode tensors
194
	for range llm.numTensor() {
195
		name, err := readGGUFString(llm, rs)
196
		if err != nil {
197
			return fmt.Errorf("failed to read tensor name: %w", err)
198
199
		}

200
201
		// dims is the number of dimensions in the tensor
		dims, err := readGGUF[uint32](llm, rs)
Bruce MacDonald's avatar
Bruce MacDonald committed
202
		if err != nil {
203
			return fmt.Errorf("failed to read tensor dimensions: %w", err)
Bruce MacDonald's avatar
Bruce MacDonald committed
204
205
		}

Michael Yang's avatar
Michael Yang committed
206
		shape := make([]uint64, dims)
207
208
		for i := 0; uint32(i) < dims; i++ {
			shape[i], err = readGGUF[uint64](llm, rs)
Bruce MacDonald's avatar
Bruce MacDonald committed
209
			if err != nil {
210
				return fmt.Errorf("failed to read tensor shape: %w", err)
Bruce MacDonald's avatar
Bruce MacDonald committed
211
212
213
			}
		}

214
		kind, err := readGGUF[uint32](llm, rs)
215
		if err != nil {
216
			return fmt.Errorf("failed to read tensor kind: %w", err)
217
218
		}

219
220
		offset, err := readGGUF[uint64](llm, rs)
		if err != nil {
221
			return fmt.Errorf("failed to read tensor offset: %w", err)
222
223
		}

224
225
		tensor := Tensor{
			Name:   name,
226
227
			Kind:   kind,
			Offset: offset,
Michael Yang's avatar
Michael Yang committed
228
			Shape:  shape[:],
Michael Yang's avatar
Michael Yang committed
229
		}
230

Michael Yang's avatar
Michael Yang committed
231
		llm.tensors = append(llm.tensors, &tensor)
232
		llm.parameters += tensor.parameters()
233
234
	}

Michael Yang's avatar
Michael Yang committed
235
236
237
	// patch KV with parameter count
	llm.kv["general.parameter_count"] = llm.parameters

Michael Yang's avatar
Michael Yang committed
238
	alignment := llm.kv.Uint("general.alignment", 32)
239

240
241
242
243
244
	offset, err := rs.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
245
	padding := ggufPadding(offset, int64(alignment))
246
247
	llm.tensorOffset = uint64(offset + padding)

Michael Yang's avatar
Michael Yang committed
248
	for _, tensor := range llm.tensors {
249
250
251
		offset, err := rs.Seek(0, io.SeekCurrent)
		if err != nil {
			return fmt.Errorf("failed to get current offset: %w", err)
252
253
		}

Michael Yang's avatar
Michael Yang committed
254
		padding := ggufPadding(offset, int64(alignment))
255
		if _, err := rs.Seek(padding, io.SeekCurrent); err != nil {
256
257
258
259
260
			return fmt.Errorf("failed to seek to init padding: %w", err)
		}

		if _, err := rs.Seek(int64(tensor.Size()), io.SeekCurrent); err != nil {
			return fmt.Errorf("failed to seek to tensor: %w", err)
Michael Yang's avatar
Michael Yang committed
261
		}
262
263
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
264
265
266
	return nil
}

267
268
269
270
func readGGUF[T any](llm *gguf, r io.Reader) (T, error) {
	var t T
	err := binary.Read(r, llm.ByteOrder, &t)
	return t, err
Bruce MacDonald's avatar
Bruce MacDonald committed
271
272
}

Michael Yang's avatar
Michael Yang committed
273
274
func writeGGUF[V any](w io.Writer, t uint32, v V) error {
	if err := binary.Write(w, binary.LittleEndian, t); err != nil {
275
276
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
277

Michael Yang's avatar
Michael Yang committed
278
	return binary.Write(w, binary.LittleEndian, v)
Bruce MacDonald's avatar
Bruce MacDonald committed
279
280
}

281
282
283
284
285
func readGGUFV1String(llm *gguf, r io.Reader) (string, error) {
	var length uint64
	if err := binary.Read(r, llm.ByteOrder, &length); err != nil {
		return "", err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
286
287

	var b bytes.Buffer
288
	if _, err := io.CopyN(&b, r, int64(length)); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
289
290
291
292
293
294
295
296
297
		return "", err
	}

	// gguf v1 strings are null-terminated
	b.Truncate(b.Len() - 1)

	return b.String(), nil
}

298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
func discardGGUFString(llm *gguf, r io.Reader) error {
	buf := llm.scratch[:8]
	_, err := io.ReadFull(r, buf)
	if err != nil {
		return err
	}

	size := int(llm.ByteOrder.Uint64(buf))
	for size > 0 {
		n, err := r.Read(llm.scratch[:min(size, cap(llm.scratch))])
		if err != nil {
			return err
		}
		size -= n
	}
	return nil
}

316
func readGGUFString(llm *gguf, r io.Reader) (string, error) {
317
	if llm.Version == 1 {
318
		return readGGUFV1String(llm, r)
319
320
	}

321
322
323
	buf := llm.scratch[:8]
	_, err := io.ReadFull(r, buf)
	if err != nil {
324
325
		return "", err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
326

327
328
329
330
331
	length := int(llm.ByteOrder.Uint64(buf))
	if length > len(llm.scratch) {
		buf = make([]byte, length)
	} else {
		buf = llm.scratch[:length]
Bruce MacDonald's avatar
Bruce MacDonald committed
332
	}
333
	clear(buf)
Bruce MacDonald's avatar
Bruce MacDonald committed
334

335
336
337
338
339
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return "", err
	}
	return string(buf), nil
Bruce MacDonald's avatar
Bruce MacDonald committed
340
341
}

Michael Yang's avatar
Michael Yang committed
342
343
func writeGGUFString(w io.Writer, s string) error {
	if err := binary.Write(w, binary.LittleEndian, ggufTypeString); err != nil {
344
345
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
346

Michael Yang's avatar
Michael Yang committed
347
	if err := binary.Write(w, binary.LittleEndian, uint64(len(s))); err != nil {
348
349
350
351
352
353
354
		return err
	}

	_, err := io.Copy(w, strings.NewReader(s))
	return err
}

355
356
357
358
359
360
361
362
363
364
type array struct {
	size   int
	values []any
}

func (a *array) MarshalJSON() ([]byte, error) {
	return json.Marshal(a.values)
}

func readGGUFV1Array(llm *gguf, r io.Reader) (*array, error) {
365
366
367
368
369
370
371
372
373
	t, err := readGGUF[uint32](llm, r)
	if err != nil {
		return nil, err
	}

	n, err := readGGUF[uint32](llm, r)
	if err != nil {
		return nil, err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
374

375
376
377
378
379
380
	a := &array{size: int(n)}
	if llm.canCollectArray(int(n)) {
		a.values = make([]any, 0, int(n))
	}

	for i := range n {
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
		var e any
		switch t {
		case ggufTypeUint8:
			e, err = readGGUF[uint8](llm, r)
		case ggufTypeInt8:
			e, err = readGGUF[int8](llm, r)
		case ggufTypeUint16:
			e, err = readGGUF[uint16](llm, r)
		case ggufTypeInt16:
			e, err = readGGUF[int16](llm, r)
		case ggufTypeUint32:
			e, err = readGGUF[uint32](llm, r)
		case ggufTypeInt32:
			e, err = readGGUF[int32](llm, r)
		case ggufTypeUint64:
			e, err = readGGUF[uint64](llm, r)
		case ggufTypeInt64:
			e, err = readGGUF[int64](llm, r)
		case ggufTypeFloat32:
			e, err = readGGUF[float32](llm, r)
		case ggufTypeFloat64:
			e, err = readGGUF[float64](llm, r)
		case ggufTypeBool:
			e, err = readGGUF[bool](llm, r)
		case ggufTypeString:
			e, err = readGGUFV1String(llm, r)
Bruce MacDonald's avatar
Bruce MacDonald committed
407
		default:
408
			return nil, fmt.Errorf("invalid array type: %d", t)
Bruce MacDonald's avatar
Bruce MacDonald committed
409
		}
410
411
412
413
		if err != nil {
			return nil, err
		}

414
415
416
		if a.values != nil {
			a.values[i] = e
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
417
418
	}

419
	return a, nil
Bruce MacDonald's avatar
Bruce MacDonald committed
420
421
}

422
func readGGUFArray(llm *gguf, r io.Reader) (*array, error) {
423
	if llm.Version == 1 {
424
		return readGGUFV1Array(llm, r)
425
426
	}

427
428
429
430
	t, err := readGGUF[uint32](llm, r)
	if err != nil {
		return nil, err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
431

432
433
434
435
	n, err := readGGUF[uint64](llm, r)
	if err != nil {
		return nil, err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
436

437
438
439
440
441
442
	a := &array{size: int(n)}
	if llm.canCollectArray(int(n)) {
		a.values = make([]any, int(n))
	}

	for i := range n {
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
		var e any
		switch t {
		case ggufTypeUint8:
			e, err = readGGUF[uint8](llm, r)
		case ggufTypeInt8:
			e, err = readGGUF[int8](llm, r)
		case ggufTypeUint16:
			e, err = readGGUF[uint16](llm, r)
		case ggufTypeInt16:
			e, err = readGGUF[int16](llm, r)
		case ggufTypeUint32:
			e, err = readGGUF[uint32](llm, r)
		case ggufTypeInt32:
			e, err = readGGUF[int32](llm, r)
		case ggufTypeUint64:
			e, err = readGGUF[uint64](llm, r)
		case ggufTypeInt64:
			e, err = readGGUF[int64](llm, r)
		case ggufTypeFloat32:
			e, err = readGGUF[float32](llm, r)
		case ggufTypeFloat64:
			e, err = readGGUF[float64](llm, r)
		case ggufTypeBool:
			e, err = readGGUF[bool](llm, r)
		case ggufTypeString:
468
469
470
471
472
			if a.values != nil {
				e, err = readGGUFString(llm, r)
			} else {
				err = discardGGUFString(llm, r)
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
473
		default:
474
			return nil, fmt.Errorf("invalid array type: %d", t)
Bruce MacDonald's avatar
Bruce MacDonald committed
475
		}
476
477
478
479
		if err != nil {
			return nil, err
		}

480
481
482
		if a.values != nil {
			a.values[i] = e
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
483
484
	}

485
	return a, nil
Bruce MacDonald's avatar
Bruce MacDonald committed
486
}
487

Michael Yang's avatar
Michael Yang committed
488
// writeGGUFArray writes a slice s of type E to the write with a gguf type of t
Michael Yang's avatar
Michael Yang committed
489
490
func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error {
	if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil {
491
492
493
		return err
	}

Michael Yang's avatar
Michael Yang committed
494
	if err := binary.Write(w, binary.LittleEndian, t); err != nil {
495
496
497
		return err
	}

Michael Yang's avatar
Michael Yang committed
498
	if err := binary.Write(w, binary.LittleEndian, uint64(len(s))); err != nil {
499
500
501
		return err
	}

Michael Yang's avatar
Michael Yang committed
502
	return binary.Write(w, binary.LittleEndian, s)
503
504
}

Michael Yang's avatar
Michael Yang committed
505
func WriteGGUF(ws io.WriteSeeker, kv KV, ts []Tensor) error {
Michael Yang's avatar
Michael Yang committed
506
507
	alignment := kv.Uint("general.alignment", 32)

Michael Yang's avatar
Michael Yang committed
508
	if err := binary.Write(ws, binary.LittleEndian, []byte("GGUF")); err != nil {
509
510
511
		return err
	}

Michael Yang's avatar
Michael Yang committed
512
	if err := binary.Write(ws, binary.LittleEndian, uint32(3)); err != nil {
513
514
515
		return err
	}

Michael Yang's avatar
Michael Yang committed
516
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(ts))); err != nil {
517
518
519
		return err
	}

Michael Yang's avatar
Michael Yang committed
520
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(kv))); err != nil {
521
522
523
		return err
	}

Michael Yang's avatar
Michael Yang committed
524
	keys := slices.Collect(maps.Keys(kv))
Michael Yang's avatar
Michael Yang committed
525
526
527
528
529
530
	slices.Sort(keys)

	for _, key := range keys {
		if err := ggufWriteKV(ws, key, kv[key]); err != nil {
			return err
		}
531
532
	}

533
534
535
536
537
538
539
	slices.SortStableFunc(ts, func(a, b Tensor) int {
		if i, j := a.block(), b.block(); i < 0 && j > 0 {
			return 1
		} else if i > 0 && j < 0 {
			return -1
		} else {
			return cmp.Compare(i, j)
540
		}
Michael Yang's avatar
Michael Yang committed
541
542
543
544
	})

	var s uint64
	for _, t := range ts {
Michael Yang's avatar
Michael Yang committed
545
		t.Offset = s + uint64(ggufPadding(int64(s), int64(alignment)))
Michael Yang's avatar
Michael Yang committed
546
		if err := ggufWriteTensorInfo(ws, t); err != nil {
547
548
			return err
		}
Michael Yang's avatar
Michael Yang committed
549
550
		s += t.Size()
	}
551

Michael Yang's avatar
Michael Yang committed
552
	for _, t := range ts {
Michael Yang's avatar
Michael Yang committed
553
		if err := ggufWriteTensor(ws, t, int64(alignment)); err != nil {
554
555
			return err
		}
Michael Yang's avatar
Michael Yang committed
556
	}
557

Michael Yang's avatar
Michael Yang committed
558
559
	return nil
}
560

Michael Yang's avatar
Michael Yang committed
561
562
563
564
func ggufWriteKV(ws io.WriteSeeker, k string, v any) error {
	slog.Debug(k, "type", fmt.Sprintf("%T", v))
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(k))); err != nil {
		return err
565
566
	}

Michael Yang's avatar
Michael Yang committed
567
568
	if err := binary.Write(ws, binary.LittleEndian, []byte(k)); err != nil {
		return err
569
570
	}

Michael Yang's avatar
Michael Yang committed
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
	var err error
	switch v := v.(type) {
	case uint32:
		err = writeGGUF(ws, ggufTypeUint32, v)
	case float32:
		err = writeGGUF(ws, ggufTypeFloat32, v)
	case bool:
		err = writeGGUF(ws, ggufTypeBool, v)
	case string:
		err = writeGGUFString(ws, v)
	case []int32:
		err = writeGGUFArray(ws, ggufTypeInt32, v)
	case []uint32:
		err = writeGGUFArray(ws, ggufTypeUint32, v)
	case []float32:
		err = writeGGUFArray(ws, ggufTypeFloat32, v)
	case []string:
		if err := binary.Write(ws, binary.LittleEndian, ggufTypeArray); err != nil {
589
590
591
			return err
		}

Michael Yang's avatar
Michael Yang committed
592
		if err := binary.Write(ws, binary.LittleEndian, ggufTypeString); err != nil {
593
594
595
			return err
		}

Michael Yang's avatar
Michael Yang committed
596
		if err := binary.Write(ws, binary.LittleEndian, uint64(len(v))); err != nil {
597
598
599
			return err
		}

Michael Yang's avatar
Michael Yang committed
600
601
		for _, e := range v {
			if err := binary.Write(ws, binary.LittleEndian, uint64(len(e))); err != nil {
602
603
604
				return err
			}

Michael Yang's avatar
Michael Yang committed
605
606
607
			if err := binary.Write(ws, binary.LittleEndian, []byte(e)); err != nil {
				return err
			}
608
		}
Michael Yang's avatar
Michael Yang committed
609
610
611
	default:
		return fmt.Errorf("improper type for '%s'", k)
	}
612

Michael Yang's avatar
Michael Yang committed
613
614
615
	return err
}

Michael Yang's avatar
Michael Yang committed
616
func ggufWriteTensorInfo(ws io.WriteSeeker, t Tensor) error {
Michael Yang's avatar
Michael Yang committed
617
618
619
	slog.Debug(t.Name, "kind", t.Kind, "shape", t.Shape, "offset", t.Offset)
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(t.Name))); err != nil {
		return err
620
621
	}

Michael Yang's avatar
Michael Yang committed
622
623
624
	if err := binary.Write(ws, binary.LittleEndian, []byte(t.Name)); err != nil {
		return err
	}
625

Michael Yang's avatar
Michael Yang committed
626
627
628
	if err := binary.Write(ws, binary.LittleEndian, uint32(len(t.Shape))); err != nil {
		return err
	}
629

Michael Yang's avatar
Michael Yang committed
630
631
	for i := range len(t.Shape) {
		if err := binary.Write(ws, binary.LittleEndian, t.Shape[len(t.Shape)-i-1]); err != nil {
632
633
			return err
		}
634
635
	}

Michael Yang's avatar
Michael Yang committed
636
637
638
639
640
641
642
	if err := binary.Write(ws, binary.LittleEndian, t.Kind); err != nil {
		return err
	}

	return binary.Write(ws, binary.LittleEndian, t.Offset)
}

Michael Yang's avatar
Michael Yang committed
643
func ggufWriteTensor(ws io.WriteSeeker, t Tensor, alignment int64) error {
Michael Yang's avatar
Michael Yang committed
644
645
646
647
648
649
650
651
652
653
654
	offset, err := ws.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
	}

	if err := binary.Write(ws, binary.LittleEndian, bytes.Repeat([]byte{0}, int(ggufPadding(offset, alignment)))); err != nil {
		return err
	}

	_, err = t.WriteTo(ws)
	return err
655
656
}

Michael Yang's avatar
Michael Yang committed
657
func ggufPadding(offset, align int64) int64 {
658
	return (align - offset%align) % align
659
}