Added BFloat16 alias.

gomlx · Jul 6, 2024 · 9502c13 · 9502c13
1 parent 4ffb520
commit 9502c13
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 7 deletions.
diff --git a/cmd/dtypes_codegen/enums.go b/cmd/dtypes_codegen/enums.go
@@ -29,6 +29,7 @@ var aliases = map[string]string{
 	"U32":     "Uint32",
 	"U64":     "Uint64",
 	"F16":     "Float16",
+	"BF16":    "BFloat16",
 	"F32":     "Float32",
 	"F64":     "Float64",
 	"C64":     "Complex64",

diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md
@@ -1,6 +1,7 @@
 # Next
 
 * Moved some `dtypes` support functionality from GoMLX to Gopjrt. 
+* Added BFloat16 alias.
 
 # v0.1.2 SuppressAbseilLoggingHack
 

diff --git a/dtypes/dtype_string.go b/dtypes/dtype_string.go
diff --git a/dtypes/gen_dtype_enum.go b/dtypes/gen_dtype_enum.go
@@ -61,11 +61,11 @@ const (
 	// Float64 is a 1:1 mapping of the corresponding C enum value defined in pjrt_c_api.h (as PJRT_Buffer_Type_F64).
 	Float64 DType = 12
 
-	// BF16 is a 1:1 mapping of the corresponding C enum value defined in pjrt_c_api.h (as PJRT_Buffer_Type_BF16).
+	// BFloat16 is a 1:1 mapping of the corresponding C enum value defined in pjrt_c_api.h (as PJRT_Buffer_Type_BF16).
 	// Truncated 16 bit floating-point format. This is similar to IEEE's 16 bit
 	// floating-point format, but uses 1 bit for the sign, 8 bits for the exponent
 	// and 7 bits for the mantissa.
-	BF16 DType = 13
+	BFloat16 DType = 13
 
 	// Complex64 is a 1:1 mapping of the corresponding C enum value defined in pjrt_c_api.h (as PJRT_Buffer_Type_C64).
 	// Complex values of fixed width.
@@ -152,6 +152,9 @@ const (
 	// F64 (or PJRT_Buffer_Type_F64) is the C enum name for Float64.
 	F64 = Float64
 
+	// BF16 (or PJRT_Buffer_Type_BF16) is the C enum name for BFloat16.
+	BF16 = BFloat16
+
 	// C64 (or PJRT_Buffer_Type_C64) is the C enum name for Complex64.
 	C64 = Complex64
 
@@ -192,7 +195,7 @@ func (dtype DType) PrimitiveType() protos.PrimitiveType {
 		return protos.PrimitiveType_F32
 	case Float64:
 		return protos.PrimitiveType_F64
-	case BF16:
+	case BFloat16:
 		return protos.PrimitiveType_BF16
 	case Complex64:
 		return protos.PrimitiveType_C64
@@ -257,7 +260,7 @@ func FromPrimitiveType(primitiveType protos.PrimitiveType) DType {
 	case protos.PrimitiveType_F64:
 		return Float64
 	case protos.PrimitiveType_BF16:
-		return BF16
+		return BFloat16
 	case protos.PrimitiveType_C64:
 		return Complex64
 	case protos.PrimitiveType_C128: