From c3cbf5c0415886747e39ba21047b5486d36d5d71 Mon Sep 17 00:00:00 2001
From: Siddhartha Bagaria <sbagaria@grailbio.com>
Date: Wed, 14 Aug 2019 17:39:18 -0700
Subject: [PATCH] Go: allow larger C array backed slices on 64 bit machines

---
 tensorflow/go/graph.go  | 7 ++++++-
 tensorflow/go/tensor.go | 7 ++++++-
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensorflow/go/graph.go b/tensorflow/go/graph.go
index 1524741fee5..785ac51cdd0 100644
--- a/tensorflow/go/graph.go
+++ b/tensorflow/go/graph.go
@@ -94,7 +94,12 @@ func (g *Graph) WriteTo(w io.Writer) (int64, error) {
 	// A []byte slice backed by C memory.
 	// See: https://github.com/golang/go/wiki/cgo#turning-c-arrays-into-go-slices
 	length := int(buf.length)
-	slice := (*[1 << 30]byte)(unsafe.Pointer(buf.data))[:length:length]
+	var slice []byte
+	if unsafe.Sizeof(unsafe.Pointer(nil)) == 8 {
+		slice = (*[1<<50 - 1]byte)(unsafe.Pointer(buf.data))[:length:length]
+	} else {
+		slice = (*[1 << 30]byte)(unsafe.Pointer(buf.data))[:length:length]
+	}
 	n, err := w.Write(slice)
 	return int64(n), err
 }
diff --git a/tensorflow/go/tensor.go b/tensorflow/go/tensor.go
index f3338f65957..9bc643ae6d2 100644
--- a/tensorflow/go/tensor.go
+++ b/tensorflow/go/tensor.go
@@ -207,7 +207,12 @@ func tensorData(c *C.TF_Tensor) []byte {
 		return nil
 	}
 	length := int(C.TF_TensorByteSize(c))
-	slice := (*[1 << 30]byte)(unsafe.Pointer(cbytes))[:length:length]
+	var slice []byte
+	if unsafe.Sizeof(unsafe.Pointer(nil)) == 8 {
+		slice = (*[1<<50 - 1]byte)(unsafe.Pointer(cbytes))[:length:length]
+	} else {
+		slice = (*[1 << 30]byte)(unsafe.Pointer(cbytes))[:length:length]
+	}
 	return slice
 }