Added XlaBuilder.Op.Builder

gomlx · Jul 13, 2024 · 2461aad · 2461aad
1 parent f62ba33
commit 2461aad
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 1 deletion.
diff --git a/cmd/xlabuilder_codegen/main.go b/cmd/xlabuilder_codegen/main.go
@@ -16,7 +16,7 @@ type OpInfo struct {
 }
 
 func main() {
-	// Read node_types.xt
+	// Read op_types.xt
 	opsInfo := make([]OpInfo, 0, 200)
 	f := must.M1(os.OpenFile(OpTypesFileName, os.O_RDONLY, os.ModePerm))
 	scanner := bufio.NewScanner(f)

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
@@ -14,6 +14,7 @@
 * Added `Buffer.Device` and `Client.NumForDevice`.
 * Properly setting client options for `pjrt.NewClient`. Added test for reading/writing `C.PJRT_NamedValues`.
 * Added `xlabuilder.Shape.Memory` and `xlabuilder.NewArrayLiteralFromAny`.
+* Added `xlabuilder.Op.Builder()`
 
 # v0.1.2 SuppressAbseilLoggingHack
 

diff --git a/xlabuilder/op.go b/xlabuilder/op.go
@@ -67,6 +67,11 @@ func newOp(opType OpType, opInputs ...*Op) *Op {
 	return op
 }
 
+// Builder returns the XlaBuilder associated with this Op.
+func (op *Op) Builder() *XlaBuilder {
+	return op.builder
+}
+
 // opFinalizer by freeing the underlying C++ resources.
 func opFinalizer(op *Op) {
 	if op.cOp == nil {