[TF-DF] [YDF] Remove old Tensorflow Status compatibility

PiperOrigin-RevId: 689818683
tensorflow · Oct 25, 2024 · 41a8f56 · 41a8f56
1 parent 932a4a3
commit 41a8f56
Show file tree

Hide file tree

Showing 16 changed files with 405 additions and 383 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Changelog
 
+## HEAD
+
+### Fix
+
+-   Some errors are now InvalidArgumentError instead of UnknownError.
+
 ## 1.10.0 - 2024-08-21
 
 ### Fix

diff --git a/tensorflow_decision_forests/keras/BUILD b/tensorflow_decision_forests/keras/BUILD
@@ -215,7 +215,6 @@ py_test(
         "@pypi_pandas//:pkg",
         "@pypi_portpicker//:pkg",
         "@release_or_nightly//:tensorflow_pkg",
-        "@release_or_nightly//:tf_keras_pkg",
         # TensorFlow /distribute:distribute_lib,
         "//tensorflow_decision_forests",
         "@ydf//yggdrasil_decision_forests/learner/distributed_gradient_boosted_trees:dgbt_py_proto",

diff --git a/tensorflow_decision_forests/keras/keras_distributed_test.py b/tensorflow_decision_forests/keras/keras_distributed_test.py
@@ -706,7 +706,7 @@ def test_in_memory_not_supported(self):
     model = tfdf.keras.DistributedGradientBoostedTreesModel(worker_logs=False)
 
     with self.assertRaisesRegex(
-        tf.errors.UnknownError,
+        tf.errors.InvalidArgumentError,
         "does not support training from in-memory datasets",
     ):
       model.fit(dataset)

diff --git a/tensorflow_decision_forests/keras/keras_test.py b/tensorflow_decision_forests/keras/keras_test.py
@@ -2923,7 +2923,8 @@ def test_node_format_blob_sequence(self):
 
   def test_check_parameters(self):
     with self.assertRaisesRegex(
-        ValueError, 'The parameter "num_trees" is smaller than the minimum'
+        tf.errors.InvalidArgumentError,
+        'The parameter "num_trees" is smaller than the minimum',
     ):
       keras.RandomForestModel(num_trees=-10)
 
@@ -3041,7 +3042,7 @@ def test_monotonic_non_compatible_learner(self):
     pd_dataset = pd.DataFrame({"f": [0, 1], "l": [0, 1]})
     tf_dataset = keras.pd_dataframe_to_tf_dataset(pd_dataset, label="l")
     with self.assertRaisesRegex(
-        tf.errors.UnknownError,
+        tf.errors.InvalidArgumentError,
         "The learner CART does not support monotonic constraints",
     ):
       model.fit(tf_dataset)
@@ -3053,7 +3054,7 @@ def test_monotonic_non_compatible_options(self):
     pd_dataset = pd.DataFrame({"f": [0, 1], "l": [0, 1]})
     tf_dataset = keras.pd_dataframe_to_tf_dataset(pd_dataset, label="l")
     with self.assertRaisesRegex(
-        tf.errors.UnknownError,
+        tf.errors.InvalidArgumentError,
         "Gradient Boosted Trees does not support monotonic constraints with"
         " use_hessian_gain=false",
     ):

diff --git a/tensorflow_decision_forests/tensorflow/ops/inference/BUILD b/tensorflow_decision_forests/tensorflow/ops/inference/BUILD
@@ -111,7 +111,11 @@ cc_library(
     name = "op",
     srcs = OP_SRCS,
     linkstatic = 1,
-    deps = OP_DEPS,
+    deps = OP_DEPS + [
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@ydf//yggdrasil_decision_forests/utils:status_macros",
+    ],
     alwayslink = 1,
 )
 
@@ -133,7 +137,6 @@ cc_library(
         "@ydf//yggdrasil_decision_forests/utils:compatibility",
         "@ydf//yggdrasil_decision_forests/utils:distribution_cc_proto",
         "@ydf//yggdrasil_decision_forests/utils:status_macros",
-        "@ydf//yggdrasil_decision_forests/utils:tensorflow",
     ] + select({
         "@ydf//yggdrasil_decision_forests:tensorflow_with_header_lib": [
             "@release_or_nightly//:tensorflow_libtensorflow_framework",