Improved support for bfloat16 dtype

Author: CoreyGiles

R/RcppExports.R

@@ -329,6 +329,10 @@ cpp_torch_float16 <- function() {
     .Call(`_torch_cpp_torch_float16`)
 }
 
+cpp_torch_bfloat16 <- function() {
+    .Call(`_torch_cpp_torch_bfloat16`)
+}
+
 cpp_torch_uint8 <- function() {
     .Call(`_torch_cpp_torch_uint8`)
 }

R/dtype.R

@@ -109,6 +109,9 @@ torch_float16 <- function() torch_dtype$new(cpp_torch_float16())
 #' @rdname torch_dtype
 #' @export
 torch_half <- function() torch_dtype$new(cpp_torch_float16())
+#' @rdname torch_dtype
+#' @export
+torch_bfloat16 <- function() torch_dtype$new(cpp_torch_bfloat16())
 
 #' @rdname torch_dtype
 #' @export

inst/include/lantern/lantern.h

@@ -256,6 +256,8 @@ LANTERN_OPTIONAL_DECLS(string_view)
   HOST_API void * lantern_Dtype_float64() {LANTERN_CHECK_LOADED void * ret = _lantern_Dtype_float64(); LANTERN_HOST_HANDLER return ret;}
   LANTERN_API void *(LANTERN_PTR _lantern_Dtype_float16)();
   HOST_API void * lantern_Dtype_float16() {LANTERN_CHECK_LOADED void * ret = _lantern_Dtype_float16(); LANTERN_HOST_HANDLER return ret;}
+  LANTERN_API void *(LANTERN_PTR _lantern_Dtype_bfloat16)();
+  HOST_API void * lantern_Dtype_bfloat16() {LANTERN_CHECK_LOADED void * ret = _lantern_Dtype_bfloat16(); LANTERN_HOST_HANDLER return ret;}
   LANTERN_API void *(LANTERN_PTR _lantern_Dtype_uint8)();
   HOST_API void * lantern_Dtype_uint8() {LANTERN_CHECK_LOADED void * ret = _lantern_Dtype_uint8(); LANTERN_HOST_HANDLER return ret;}
   LANTERN_API void *(LANTERN_PTR _lantern_Dtype_int8)();
@@ -10530,6 +10532,7 @@ bool lanternInit(const std::string &libPath, std::string *pError)
   LOAD_SYMBOL(_lantern_Dtype_float32);
   LOAD_SYMBOL(_lantern_Dtype_float64);
   LOAD_SYMBOL(_lantern_Dtype_float16);
+  LOAD_SYMBOL(_lantern_Dtype_bfloat16);
   LOAD_SYMBOL(_lantern_Dtype_uint8);
   LOAD_SYMBOL(_lantern_Dtype_int8);
   LOAD_SYMBOL(_lantern_Dtype_int16);

src/RcppExports.cpp

@@ -910,6 +910,16 @@ BEGIN_RCPP
     return rcpp_result_gen;
 END_RCPP
 }
+// cpp_torch_bfloat16
+Rcpp::XPtr<XPtrTorchDtype> cpp_torch_bfloat16();
+RcppExport SEXP _torch_cpp_torch_bfloat16() {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    rcpp_result_gen = Rcpp::wrap(cpp_torch_bfloat16());
+    return rcpp_result_gen;
+END_RCPP
+}
 // cpp_torch_uint8
 Rcpp::XPtr<XPtrTorchDtype> cpp_torch_uint8();
 RcppExport SEXP _torch_cpp_torch_uint8() {
@@ -48231,6 +48241,7 @@ static const R_CallMethodDef CallEntries[] = {
     {"_torch_cpp_torch_float32", (DL_FUNC) &_torch_cpp_torch_float32, 0},
     {"_torch_cpp_torch_float64", (DL_FUNC) &_torch_cpp_torch_float64, 0},
     {"_torch_cpp_torch_float16", (DL_FUNC) &_torch_cpp_torch_float16, 0},
+    {"_torch_cpp_torch_bfloat16", (DL_FUNC) &_torch_cpp_torch_bfloat16, 0},
     {"_torch_cpp_torch_uint8", (DL_FUNC) &_torch_cpp_torch_uint8, 0},
     {"_torch_cpp_torch_int8", (DL_FUNC) &_torch_cpp_torch_int8, 0},
     {"_torch_cpp_torch_int16", (DL_FUNC) &_torch_cpp_torch_int16, 0},

src/dtype.cpp

@@ -20,6 +20,11 @@ Rcpp::XPtr<XPtrTorchDtype> cpp_torch_float16() {
   return make_xptr<XPtrTorchDtype>(lantern_Dtype_float16());
 }
 
+// [[Rcpp::export]]
+Rcpp::XPtr<XPtrTorchDtype> cpp_torch_bfloat16() {
+  return make_xptr<XPtrTorchDtype>(lantern_Dtype_bfloat16());
+}
+
 // [[Rcpp::export]]
 Rcpp::XPtr<XPtrTorchDtype> cpp_torch_uint8() {
   return make_xptr<XPtrTorchDtype>(lantern_Dtype_uint8());

src/lantern/include/lantern/lantern.h

@@ -256,6 +256,8 @@ LANTERN_OPTIONAL_DECLS(string_view)
   HOST_API void * lantern_Dtype_float64() {LANTERN_CHECK_LOADED void * ret = _lantern_Dtype_float64(); LANTERN_HOST_HANDLER return ret;}
   LANTERN_API void *(LANTERN_PTR _lantern_Dtype_float16)();
   HOST_API void * lantern_Dtype_float16() {LANTERN_CHECK_LOADED void * ret = _lantern_Dtype_float16(); LANTERN_HOST_HANDLER return ret;}
+  LANTERN_API void *(LANTERN_PTR _lantern_Dtype_bfloat16)();
+  HOST_API void * lantern_Dtype_bfloat16() {LANTERN_CHECK_LOADED void * ret = _lantern_Dtype_bfloat16(); LANTERN_HOST_HANDLER return ret;}
   LANTERN_API void *(LANTERN_PTR _lantern_Dtype_uint8)();
   HOST_API void * lantern_Dtype_uint8() {LANTERN_CHECK_LOADED void * ret = _lantern_Dtype_uint8(); LANTERN_HOST_HANDLER return ret;}
   LANTERN_API void *(LANTERN_PTR _lantern_Dtype_int8)();
@@ -10530,6 +10532,7 @@ bool lanternInit(const std::string &libPath, std::string *pError)
   LOAD_SYMBOL(_lantern_Dtype_float32);
   LOAD_SYMBOL(_lantern_Dtype_float64);
   LOAD_SYMBOL(_lantern_Dtype_float16);
+  LOAD_SYMBOL(_lantern_Dtype_bfloat16);
   LOAD_SYMBOL(_lantern_Dtype_uint8);
   LOAD_SYMBOL(_lantern_Dtype_int8);
   LOAD_SYMBOL(_lantern_Dtype_int16);

src/lantern/src/Dtype.cpp

@@ -13,6 +13,7 @@
 LANTERN_DTYPE_FUN(float16, kFloat16)
 LANTERN_DTYPE_FUN(float32, kFloat32)
 LANTERN_DTYPE_FUN(float64, kFloat64)
+LANTERN_DTYPE_FUN(bfloat16, kBFloat16)
 LANTERN_DTYPE_FUN(int8, kInt8)
 LANTERN_DTYPE_FUN(int16, kInt16)
 LANTERN_DTYPE_FUN(int32, kInt32)