Add T=bfloat16 to custom_ops registration (#2688)

tensorflow · Apr 12, 2022 · c585796 · c585796
1 parent 50530e8
commit c585796
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 10 deletions.
diff --git a/tensorflow_addons/custom_ops/image/cc/ops/distort_image_ops.cc b/tensorflow_addons/custom_ops/image/cc/ops/distort_image_ops.cc
@@ -31,7 +31,8 @@ REGISTER_OP("Addons>AdjustHsvInYiq")
     .Input("scale_s: float")
     .Input("scale_v: float")
     .Output("output: T")
-    .Attr("T: {uint8, int8, int16, int32, int64, half, float, double}")
+    .Attr(
+        "T: {uint8, int8, int16, int32, int64, half, float, double, bfloat16}")
     .SetShapeFn([](InferenceContext* c) {
       ShapeHandle images, delta_h, scale_s, scale_v;
 
@@ -70,4 +71,4 @@ output: The hsv-adjusted image or images. No clipping will be done in this op.
 )Doc");
 
 }  // end namespace addons
-}  // namespace tensorflow
+}  // namespace tensorflow
diff --git a/tensorflow_addons/custom_ops/image/cc/ops/image_ops.cc b/tensorflow_addons/custom_ops/image/cc/ops/image_ops.cc
@@ -55,7 +55,7 @@ components: Component ids for each pixel in "image". Same shape as "image". Zero
 
 REGISTER_OP("Addons>EuclideanDistanceTransform")
     .Input("images: uint8")
-    .Attr("dtype: {float16, float32, float64}")
+    .Attr("dtype: {bfloat16, float16, float32, float64}")
     .Output("transformed_images: dtype")
     .SetShapeFn(shape_inference::UnchangedShape)
     .Doc(EuclideanDistanceTransformDoc);
@@ -65,9 +65,9 @@ REGISTER_OP("Addons>ImageConnectedComponents")
     .Output("components: int64")
     .Attr(
         "dtype: {int64, int32, uint16, int16, uint8, int8, half, float, "
-        "double, bool, string}")
+        "bfloat16, double, bool, string}")
     .SetShapeFn(shape_inference::UnchangedShape)
     .Doc(ImageConnectedComponentsDoc);
 
 }  // end namespace addons
-}  // namespace tensorflow
+}  // namespace tensorflow
diff --git a/tensorflow_addons/custom_ops/image/cc/ops/resampler_ops.cc b/tensorflow_addons/custom_ops/image/cc/ops/resampler_ops.cc
@@ -29,7 +29,7 @@ REGISTER_OP("Addons>Resampler")
     .Input("data: T")
     .Input("warp: T")
     .Output("output: T")
-    .Attr("T: {half, float, double}")
+    .Attr("T: {bfloat16, half, float, double}")
     .SetShapeFn([](InferenceContext* c) {
       ShapeHandle data;
       ShapeHandle warp;
@@ -53,7 +53,7 @@ REGISTER_OP("Addons>ResamplerGrad")
     .Input("grad_output: T")
     .Output("grad_data: T")
     .Output("grad_warp: T")
-    .Attr("T: {half, float, double}")
+    .Attr("T: {bfloat16, half, float, double}")
     .SetShapeFn([](InferenceContext* c) {
       c->set_output(0, c->input(0));
       c->set_output(1, c->input(1));
@@ -62,4 +62,4 @@ REGISTER_OP("Addons>ResamplerGrad")
     .Doc(R"doc(Resampler Grad op.)doc");
 
 }  // namespace addons
-}  // namespace tensorflow
+}  // namespace tensorflow
diff --git a/tensorflow_addons/custom_ops/layers/cc/ops/embedding_bag_ops.cc b/tensorflow_addons/custom_ops/layers/cc/ops/embedding_bag_ops.cc
@@ -28,7 +28,7 @@ REGISTER_OP("Addons>EmbeddingBag")
     .Input("params: T")
     .Input("weights: T")
     .Output("output: T")
-    .Attr("T: {half, float, double}")
+    .Attr("T: {bfloat16, half, float, double}")
     .Attr("Tindices: {int32, int64}")
     .Attr("combiner: {'SUM', 'MEAN'} = 'SUM'")
     .SetShapeFn([](InferenceContext* c) {
@@ -51,7 +51,7 @@ REGISTER_OP("Addons>EmbeddingBagGrad")
     .Input("grads: T")
     .Output("params_grads: T")
     .Output("weights_grads: T")
-    .Attr("T: {half, float, double}")
+    .Attr("T: {bfloat16, half, float, double}")
     .Attr("Tindices: {int32, int64}")
     .Attr("combiner: {'SUM', 'MEAN'} = 'SUM'")
     .SetShapeFn([](InferenceContext* c) {