From 7961b3eca428bf895b158152e596e0d220f6019a Mon Sep 17 00:00:00 2001 From: Jonas Knobloch Date: Fri, 17 Apr 2026 14:51:26 +0200 Subject: Add onnx module --- onnx/cmd/extract/main.go | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 onnx/cmd/extract/main.go (limited to 'onnx/cmd/extract/main.go') diff --git a/onnx/cmd/extract/main.go b/onnx/cmd/extract/main.go new file mode 100644 index 0000000..d8f8358 --- /dev/null +++ b/onnx/cmd/extract/main.go @@ -0,0 +1,23 @@ +package main + +import ( + "fmt" + "log" + + "go.jknobloc.com/x/onnx" + "go.jknobloc.com/x/tensor" +) + +func main() { + data, shape, err := onnx.ExtractInitializer("gpt2/models/base/model.onnx", "transformer.wte.weight") + + if err != nil { + log.Fatal(err) + } + + t := tensor.NewDense[float32](shape, data) + + eot := t.Select(0, 50256).Contiguous() + + fmt.Println(eot) +} -- cgit v1.3.1