Updated to use new gotype package.

Author: janpfeifer

pjrt/buffers.go

@@ -14,6 +14,7 @@ import (
 	"unsafe"
 
 	"github.com/gomlx/compute/dtypes"
+	"github.com/gomlx/compute/dtypes/gotype"
 	"github.com/gomlx/go-xla/internal/shapeinference"
 	"github.com/gomlx/go-xla/types/shapes"
 	"github.com/pkg/errors"
@@ -222,7 +223,7 @@ func (b *Buffer) Client() *Client {
 }
 
 // ScalarToRaw generates the raw values needed by BufferFromHostConfig.FromRawData to feed a simple scalar value.
-func ScalarToRaw[T dtypes.Supported](value T) ([]byte, dtypes.DType, []int) {
+func ScalarToRaw[T gotype.Supported](value T) ([]byte, dtypes.DType, []int) {
 	dtype := dtypes.FromGenericsType[T]()
 	rawSlice := unsafe.Slice((*byte)(unsafe.Pointer(&value)), int(unsafe.Sizeof(value)))
 	return rawSlice, dtype, nil // empty dimensions for scalar
@@ -252,7 +253,7 @@ func (b *Buffer) Size() (int, error) {
 }
 
 // BufferToScalar is a generic function that transfer a Buffer back to host as a scalar of the given type.
-func BufferToScalar[T dtypes.Supported](b *Buffer) (value T, err error) {
+func BufferToScalar[T gotype.Supported](b *Buffer) (value T, err error) {
 	dst := unsafe.Slice((*byte)(unsafe.Pointer(&value)), unsafe.Sizeof(value))
 	err = b.ToHost(dst)
 	return
@@ -262,7 +263,7 @@ func BufferToScalar[T dtypes.Supported](b *Buffer) (value T, err error) {
 //
 // It is a shortcut to Client.BufferFromHost call with default parameters.
 // If you need more control where the value will be used you'll have to use Client.BufferFromHost instead.
-func ScalarToBuffer[T dtypes.Supported](client *Client, value T) (b *Buffer, err error) {
+func ScalarToBuffer[T gotype.Supported](client *Client, value T) (b *Buffer, err error) {
 	dtype := dtypes.FromGenericsType[T]()
 	src := unsafe.Slice((*byte)(unsafe.Pointer(&value)), unsafe.Sizeof(value))
 	return client.BufferFromHost().FromRawData(src, dtype, nil).Done()
@@ -272,7 +273,7 @@ func ScalarToBuffer[T dtypes.Supported](client *Client, value T) (b *Buffer, err
 //
 // It is a shortcut to Client.BufferFromHost call with default parameters.
 // If you need more control where the value will be used you'll have to use Client.BufferFromHost instead.
-func ScalarToBufferOnDeviceNum[T dtypes.Supported](client *Client, deviceNum int, value T) (b *Buffer, err error) {
+func ScalarToBufferOnDeviceNum[T gotype.Supported](client *Client, deviceNum int, value T) (b *Buffer, err error) {
 	dtype := dtypes.FromGenericsType[T]()
 	src := unsafe.Slice((*byte)(unsafe.Pointer(&value)), unsafe.Sizeof(value))
 	return client.BufferFromHost().FromRawData(src, dtype, nil).ToDeviceNum(deviceNum).Done()
@@ -283,15 +284,15 @@ func ScalarToBufferOnDeviceNum[T dtypes.Supported](client *Client, deviceNum int
 //
 // It is a shortcut to Client.BufferFromHost call with default parameters.
 // If you need more control where the value will be used you'll have to use Client.BufferFromHost instead.
-func ArrayToBuffer[T dtypes.Supported](client *Client, flatValues []T, dimensions ...int) (b *Buffer, err error) {
+func ArrayToBuffer[T gotype.Supported](client *Client, flatValues []T, dimensions ...int) (b *Buffer, err error) {
 	if len(dimensions) == 0 && len(flatValues) != 1 {
 		return nil, errors.Errorf("ArrayToBuffer not given any dimensions (indicating a scalar), but len(flatValues) == %d", len(flatValues))
 	}
 	return client.BufferFromHost().FromFlatDataWithDimensions(flatValues, dimensions).Done()
 }
 
 // BufferToArray transfers the buffer to an array defined by a slice with its flat values, and returns also its underlying dimensions.
-func BufferToArray[T dtypes.Supported](buffer *Buffer) (flatValues []T, dimensions []int, err error) {
+func BufferToArray[T gotype.Supported](buffer *Buffer) (flatValues []T, dimensions []int, err error) {
 	if err = buffer.Check(); err != nil {
 		return
 	}

pjrt/pjrt_test.go

@@ -10,7 +10,7 @@ import (
 
 	"testing"
 
-	"github.com/gomlx/compute/dtypes"
+	"github.com/gomlx/compute/dtypes/gotype"
 	"github.com/pkg/errors"
 	"k8s.io/klog/v2"
 )
@@ -88,7 +88,7 @@ func compile(t *testing.T, client *Client, comp XlaComputation) (exec *LoadedExe
 // execWithScalars executes the program on the input value given, and return the output.
 // Both input and output expected to be a scalar.
 // Any errors fail the test.
-func execWithScalars[T dtypes.Supported](t *testing.T, client *Client, exec *LoadedExecutable, input T) T {
+func execWithScalars[T gotype.Supported](t *testing.T, client *Client, exec *LoadedExecutable, input T) T {
 	inputBuffer, err := ScalarToBuffer(client, input)
 	requireNoError(t, err, "Failed to create on-device buffer for input %v", input)
 	defer func() { requireNoError(t, inputBuffer.Destroy()) }()
@@ -105,7 +105,7 @@ func execWithScalars[T dtypes.Supported](t *testing.T, client *Client, exec *Loa
 	return output
 }
 
-func execWithSlices[T dtypes.Supported](t *testing.T, client *Client, exec *LoadedExecutable, input []T) (flat []T, dims []int) {
+func execWithSlices[T gotype.Supported](t *testing.T, client *Client, exec *LoadedExecutable, input []T) (flat []T, dims []int) {
 	inputBuffer, err := ArrayToBuffer(client, input, len(input))
 	requireNoError(t, err, "Failed to create on-device buffer for input %v", input)
 	defer func() { requireNoError(t, inputBuffer.Destroy()) }()
@@ -122,7 +122,7 @@ func execWithSlices[T dtypes.Supported](t *testing.T, client *Client, exec *Load
 	return
 }
 
-func execArrayOutput[T dtypes.Supported](t *testing.T, client *Client, exec *LoadedExecutable) (flat []T, dims []int) {
+func execArrayOutput[T gotype.Supported](t *testing.T, client *Client, exec *LoadedExecutable) (flat []T, dims []int) {
 	outputBuffers := capture(exec.Execute().Done()).Test(t)
 	assertLen(t, outputBuffers, 1, "Expected only one output")
 	defer func() { requireNoError(t, outputBuffers[0].Destroy()) }()
@@ -137,7 +137,7 @@ func execArrayOutput[T dtypes.Supported](t *testing.T, client *Client, exec *Loa
 }
 
 // execScalarOutput executes the LoadedExecutable with no inputs, and a scalar output of the given type.
-func execScalarOutput[T dtypes.Supported](t *testing.T, client *Client, exec *LoadedExecutable) (value T) {
+func execScalarOutput[T gotype.Supported](t *testing.T, client *Client, exec *LoadedExecutable) (value T) {
 	outputBuffers := capture(exec.Execute().Done()).Test(t)
 	assertLen(t, outputBuffers, 1, "Expected only one output")
 	defer func() { requireNoError(t, outputBuffers[0].Destroy()) }()

types/shapes/dtype.go

@@ -23,6 +23,7 @@ import (
 	"github.com/gomlx/compute/dtypes"
 	"github.com/gomlx/compute/dtypes/bfloat16"
 	"github.com/gomlx/compute/dtypes/float16"
+	"github.com/gomlx/compute/dtypes/gotype"
 	"github.com/pkg/errors"
 )
 
@@ -32,7 +33,7 @@ import (
 // It doesn't work for if T (the output type) is a complex number.
 // If value is a complex number, it converts by taking the real part of the number and
 // discarding the imaginary part.
-func ConvertTo[T dtypes.NumberNotComplex](value any) T {
+func ConvertTo[T gotype.NumericNotComplex](value any) T {
 	t, ok := value.(T)
 	if ok {
 		return t

types/shapes/shape.go

@@ -82,6 +82,7 @@ import (
 	"strings"
 
 	"github.com/gomlx/compute/dtypes"
+	"github.com/gomlx/compute/dtypes/gotype"
 	"github.com/pkg/errors"
 )
 
@@ -147,7 +148,7 @@ func Make(dtype dtypes.DType, dimensions ...int) Shape {
 }
 
 // Scalar returns a scalar Shape for the given type.
-func Scalar[T dtypes.Number]() Shape {
+func Scalar[T gotype.Numeric]() Shape {
 	return Shape{DType: dtypes.FromGenericsType[T]()}
 }