gguf.go 13.5 KB
Newer Older
Bruce MacDonald's avatar
Bruce MacDonald committed
1
2
3
4
package llm

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
11
	"log/slog"
	"slices"
12
	"strings"
Michael Yang's avatar
Michael Yang committed
13
14

	"golang.org/x/exp/maps"
Bruce MacDonald's avatar
Bruce MacDonald committed
15
16
)

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

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

	V1 struct {
		NumTensor uint32
		NumKV     uint32
	}

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

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

	maxArraySize int
}

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

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

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

53
	var err error
Bruce MacDonald's avatar
Bruce MacDonald committed
54
55
	switch c.Version {
	case 1:
56
57
58
		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
59
	default:
60
61
62
63
		err = binary.Read(rs, c.ByteOrder, &c.V3)
	}
	if err != nil {
		return nil, err
Bruce MacDonald's avatar
Bruce MacDonald committed
64
65
	}

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

	return model, nil
}

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

90
91
type gguf struct {
	*containerGGUF
92

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

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

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

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

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

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

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

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

142
143
144
145
146
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 {
147
148
149
			return err
		}

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

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
184
		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)
185
186
		}

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

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

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

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

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

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

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

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

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

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

	alignment, ok := llm.kv["general.alignment"].(uint32)
240
241
242
	if !ok {
		alignment = 32
	}
243

244
245
246
247
248
	offset, err := rs.Seek(0, io.SeekCurrent)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
249
	padding := ggufPadding(offset, int64(alignment))
250
251
	llm.tensorOffset = uint64(offset + padding)

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

Michael Yang's avatar
Michael Yang committed
258
		padding := ggufPadding(offset, int64(alignment))
259
		if _, err := rs.Seek(padding, io.SeekCurrent); err != nil {
260
261
262
263
264
			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
265
		}
266
267
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
268
269
270
	return nil
}

271
272
273
274
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
275
276
}

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

Michael Yang's avatar
Michael Yang committed
282
	return binary.Write(w, binary.LittleEndian, v)
Bruce MacDonald's avatar
Bruce MacDonald committed
283
284
}

285
286
287
288
289
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
290
291

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

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

	return b.String(), nil
}

302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
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
}

320
func readGGUFString(llm *gguf, r io.Reader) (string, error) {
321
	if llm.Version == 1 {
322
		return readGGUFV1String(llm, r)
323
324
	}

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

331
332
333
334
335
	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
336
	}
337
	clear(buf)
Bruce MacDonald's avatar
Bruce MacDonald committed
338

339
340
341
342
343
	_, err = io.ReadFull(r, buf)
	if err != nil {
		return "", err
	}
	return string(buf), nil
Bruce MacDonald's avatar
Bruce MacDonald committed
344
345
}

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

Michael Yang's avatar
Michael Yang committed
351
	if err := binary.Write(w, binary.LittleEndian, uint64(len(s))); err != nil {
352
353
354
355
356
357
358
		return err
	}

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

359
360
361
362
363
364
365
366
367
368
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) {
369
370
371
372
373
374
375
376
377
	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
378

379
380
381
382
383
384
	a := &array{size: int(n)}
	if llm.canCollectArray(int(n)) {
		a.values = make([]any, 0, int(n))
	}

	for i := range n {
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
		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
411
		default:
412
			return nil, fmt.Errorf("invalid array type: %d", t)
Bruce MacDonald's avatar
Bruce MacDonald committed
413
		}
414
415
416
417
		if err != nil {
			return nil, err
		}

418
419
420
		if a.values != nil {
			a.values[i] = e
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
421
422
	}

423
	return a, nil
Bruce MacDonald's avatar
Bruce MacDonald committed
424
425
}

426
func readGGUFArray(llm *gguf, r io.Reader) (*array, error) {
427
	if llm.Version == 1 {
428
		return readGGUFV1Array(llm, r)
429
430
	}

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

436
437
438
439
	n, err := readGGUF[uint64](llm, r)
	if err != nil {
		return nil, err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
440

441
442
443
444
445
446
	a := &array{size: int(n)}
	if llm.canCollectArray(int(n)) {
		a.values = make([]any, int(n))
	}

	for i := range n {
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
		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:
472
473
474
475
476
			if a.values != nil {
				e, err = readGGUFString(llm, r)
			} else {
				err = discardGGUFString(llm, r)
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
477
		default:
478
			return nil, fmt.Errorf("invalid array type: %d", t)
Bruce MacDonald's avatar
Bruce MacDonald committed
479
		}
480
481
482
483
		if err != nil {
			return nil, err
		}

484
485
486
		if a.values != nil {
			a.values[i] = e
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
487
488
	}

489
	return a, nil
Bruce MacDonald's avatar
Bruce MacDonald committed
490
}
491

Michael Yang's avatar
Michael Yang committed
492
493
func writeGGUFArray[S ~[]E, E any](w io.Writer, t uint32, s S) error {
	if err := binary.Write(w, binary.LittleEndian, ggufTypeArray); err != nil {
494
495
496
		return err
	}

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

Michael Yang's avatar
Michael Yang committed
501
	if err := binary.Write(w, binary.LittleEndian, uint64(len(s))); err != nil {
502
503
504
505
		return err
	}

	for _, e := range s {
Michael Yang's avatar
Michael Yang committed
506
		if err := binary.Write(w, binary.LittleEndian, e); err != nil {
507
508
509
510
511
512
513
			return err
		}
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
514
515
func WriteGGUF(ws io.WriteSeeker, kv KV, ts []*Tensor) error {
	if err := binary.Write(ws, binary.LittleEndian, []byte("GGUF")); err != nil {
516
517
518
		return err
	}

Michael Yang's avatar
Michael Yang committed
519
	if err := binary.Write(ws, binary.LittleEndian, uint32(3)); err != nil {
520
521
522
		return err
	}

Michael Yang's avatar
Michael Yang committed
523
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(ts))); err != nil {
524
525
526
		return err
	}

Michael Yang's avatar
Michael Yang committed
527
	if err := binary.Write(ws, binary.LittleEndian, uint64(len(kv))); err != nil {
528
529
530
		return err
	}

Michael Yang's avatar
Michael Yang committed
531
532
533
534
535
536
537
	keys := maps.Keys(kv)
	slices.Sort(keys)

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

Michael Yang's avatar
Michael Yang committed
540
541
542
543
544
545
	slices.SortFunc(ts, func(a, b *Tensor) int {
		var i, j int
		if n, err := fmt.Sscanf(a.Name, "blk.%d", &i); err != nil || n != 1 {
			return cmp.Compare(a.Name, b.Name)
		} else if n, err := fmt.Sscanf(b.Name, "blk.%d", &j); err != nil || n != 1 {
			return cmp.Compare(a.Name, b.Name)
546
547
		}

Michael Yang's avatar
Michael Yang committed
548
549
550
551
552
553
554
		return cmp.Compare(i, j)
	})

	var s uint64
	for _, t := range ts {
		t.Offset = s
		if err := ggufWriteTensorInfo(ws, t); err != nil {
555
556
			return err
		}
Michael Yang's avatar
Michael Yang committed
557
558
		s += t.Size()
	}
559

Michael Yang's avatar
Michael Yang committed
560
561
562
	var alignment int64 = 32
	for _, t := range ts {
		if err := ggufWriteTensor(ws, t, alignment); err != nil {
563
564
			return err
		}
Michael Yang's avatar
Michael Yang committed
565
	}
566

Michael Yang's avatar
Michael Yang committed
567
568
	return nil
}
569

Michael Yang's avatar
Michael Yang committed
570
571
572
573
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
574
575
	}

Michael Yang's avatar
Michael Yang committed
576
577
	if err := binary.Write(ws, binary.LittleEndian, []byte(k)); err != nil {
		return err
578
579
	}

Michael Yang's avatar
Michael Yang committed
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
	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 {
598
599
600
			return err
		}

Michael Yang's avatar
Michael Yang committed
601
		if err := binary.Write(ws, binary.LittleEndian, ggufTypeString); err != nil {
602
603
604
			return err
		}

Michael Yang's avatar
Michael Yang committed
605
		if err := binary.Write(ws, binary.LittleEndian, uint64(len(v))); err != nil {
606
607
608
			return err
		}

Michael Yang's avatar
Michael Yang committed
609
610
		for _, e := range v {
			if err := binary.Write(ws, binary.LittleEndian, uint64(len(e))); err != nil {
611
612
613
				return err
			}

Michael Yang's avatar
Michael Yang committed
614
615
616
			if err := binary.Write(ws, binary.LittleEndian, []byte(e)); err != nil {
				return err
			}
617
		}
Michael Yang's avatar
Michael Yang committed
618
619
620
	default:
		return fmt.Errorf("improper type for '%s'", k)
	}
621

Michael Yang's avatar
Michael Yang committed
622
623
624
625
626
627
628
	return err
}

func ggufWriteTensorInfo(ws io.WriteSeeker, t *Tensor) error {
	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
629
630
	}

Michael Yang's avatar
Michael Yang committed
631
632
633
	if err := binary.Write(ws, binary.LittleEndian, []byte(t.Name)); err != nil {
		return err
	}
634

Michael Yang's avatar
Michael Yang committed
635
636
637
	if err := binary.Write(ws, binary.LittleEndian, uint32(len(t.Shape))); err != nil {
		return err
	}
638

Michael Yang's avatar
Michael Yang committed
639
640
	for i := range len(t.Shape) {
		if err := binary.Write(ws, binary.LittleEndian, t.Shape[len(t.Shape)-i-1]); err != nil {
641
642
			return err
		}
643
644
	}

Michael Yang's avatar
Michael Yang committed
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
	if err := binary.Write(ws, binary.LittleEndian, t.Kind); err != nil {
		return err
	}

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

func ggufWriteTensor(ws io.WriteSeeker, t *Tensor, alignment int64) error {
	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
664
665
}

Michael Yang's avatar
Michael Yang committed
666
func ggufPadding(offset, align int64) int64 {
667
	return (align - offset%align) % align
668
}