Merge pull request #33 from gomlx/dotgeneral-api-fix

Fix DotGeneral output dtype and update compute dependency

Author: janpfeifer

compute/xla/function.go

@@ -43,6 +43,24 @@ func (f *Function) Parent() compute.Function {
 	return f.parent
 }
 
+// Builder returns the builder for this function.
+func (f *Function) Builder() compute.Builder {
+	return f.builder
+}
+
+// Shape returns the shape of a value in the function.
+func (f *Function) Shape(v compute.Value) (shapes.Shape, error) {
+	var s shapes.Shape
+	n, ok := v.(*Node)
+	if !ok {
+		return s, errors.Errorf("value is not a Node for a XLA backend, instead got a %T", v)
+	}
+	if err := f.CheckValid(); err != nil {
+		return s, err
+	}
+	return n.shape, nil
+}
+
 // Closure creates a new closure function, within this function.
 func (f *Function) Closure() (compute.Function, error) {
 	if err := f.CheckValid(); err != nil {

compute/xla/ops_dotgeneral.go

@@ -23,8 +23,10 @@ import (
 //   - Contracted (contracting axes), where the output does multiply the values and reduce sum
 //     those dimensions.
 //
-// It follows that the resulting dimension number starts with the batch dimension, then the 'lhs'
-// non-contracting/non-batch dimension, and finally the 'rhs' non-contracting/non-batch dimension.
+// The resulting shape is [batchIndices..., <lhs cross indices...>, <rhs cross indices...>], the
+// indices come in the order they were provided. The output dtype is by default the same as
+// the input, except if configured otherwise in config.OutputDType.
+//
 // It provides the basic means of implementing Einsum.
 //
 // The XLA implementation only supports accumulation in F32 (if different than the input dtypes).
@@ -39,11 +41,14 @@ func (f *Function) DotGeneral(
 	}
 	lhsNode := nodes[0]
 	rhsNode := nodes[1]
-	dtype := lhsNode.shape.DType
-	accumulationDType := dtype
+	inputDType := lhsNode.shape.DType
+	accumulationDType := inputDType
 	if config.AccumulatorDType != dtypes.InvalidDType {
 		accumulationDType = config.AccumulatorDType
 	}
+	if config.OutputDType == dtypes.InvalidDType {
+		config.OutputDType = inputDType
+	}
 
 	// Introduce a dependency to create a schedule barrier to prevent the compiler to try
 	// to start processing weights too early (like converting them to the accumulationDType).
@@ -55,14 +60,14 @@ func (f *Function) DotGeneral(
 	// The CUDA backend does not seem to need this (likely they do something like that internally).
 	isCPU := f.builder.backend.plugin.IsCPU()
 	lhsReady, rhsReady := lhsNode, rhsNode
-	if accumulationDType != dtype && isCPU {
+	if accumulationDType != inputDType && isCPU {
 		lhsReady, rhsReady, err = dotGeneralAddDependency(f, lhsReady, rhsReady)
 		if err != nil {
 			return nil, err
 		}
 	}
 
-	if accumulationDType != dtype {
+	if accumulationDType != inputDType {
 		if accumulationDType != dtypes.F32 {
 			// XLA only supports accumulation in F32 (if different than the input dtypes).
 			// For other accumulation dtypes, we convert the inputs to the type.
@@ -76,7 +81,7 @@ func (f *Function) DotGeneral(
 				return nil, errors.WithMessagef(err, "failed to convert rhs to accumulation dtype")
 			}
 			lhsReady, rhsReady = lhsReadyValue.(*Node), rhsReadyValue.(*Node)
-			dtype = accumulationDType
+			inputDType = accumulationDType
 		}
 	}
 
@@ -89,7 +94,7 @@ func (f *Function) DotGeneral(
 	//
 	// So we make a simplification here: we only set the "dot_algorithm" is the plugin is not CUDA.
 	useTF32 := accumulationDType == dtypes.F32 && f.builder.backend.DotGeneralUseTF32
-	if useTF32 || accumulationDType != dtype {
+	if useTF32 || accumulationDType != inputDType {
 		isCUDA := f.builder.backend.plugin.IsCUDA()
 		if isCUDA {
 			// Set "precision_config"

go.mod

@@ -4,7 +4,7 @@ go 1.26
 
 require (
 	github.com/gofrs/flock v0.13.0
-	github.com/gomlx/compute v0.0.0-20260502161841-2635c422f98b
+	github.com/gomlx/compute v0.0.0-20260505135155-9f8cc2c04803
 	github.com/janpfeifer/go-benchmarks v0.1.1
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.11.1

go.sum

@@ -7,8 +7,8 @@ github.com/go-logr/logr v1.4.3 h1:CjnDlHq8ikf6E492q6eKboGOC0T8CDaOvkHCIg8idEI=
 github.com/go-logr/logr v1.4.3/go.mod h1:9T104GzyrTigFIr8wt5mBrctHMim0Nb2HLGrmQ40KvY=
 github.com/gofrs/flock v0.13.0 h1:95JolYOvGMqeH31+FC7D2+uULf6mG61mEZ/A8dRYMzw=
 github.com/gofrs/flock v0.13.0/go.mod h1:jxeyy9R1auM5S6JYDBhDt+E2TCo7DkratH4Pgi8P+Z0=
-github.com/gomlx/compute v0.0.0-20260502161841-2635c422f98b h1:eI6OSkBrZlBIJFWhnMQQvDiYYxmhKEOoPE48F2xj89k=
-github.com/gomlx/compute v0.0.0-20260502161841-2635c422f98b/go.mod h1:MDTT683Wvq7IMSKXDGLvV3Co7HDpx/Xsgp2RbF5lhNM=
+github.com/gomlx/compute v0.0.0-20260505135155-9f8cc2c04803 h1:9wh+H/ui5eD4ah4tA7vRTN/FiGljh8tOelKjpm1gbV4=
+github.com/gomlx/compute v0.0.0-20260505135155-9f8cc2c04803/go.mod h1:MDTT683Wvq7IMSKXDGLvV3Co7HDpx/Xsgp2RbF5lhNM=
 github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
 github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
 github.com/janpfeifer/go-benchmarks v0.1.1 h1:gLLy07/JrOKSnMWeUxSnjTdhkglgmrNR2IBDnR4kRqw=