Add MAE loss test.

PiperOrigin-RevId: 574830522

CHANGELOG.md

@@ -5,6 +5,8 @@
 ### Features
 
 -   Add support for monotonic constraints.
+-   Add support for Mean average error (MAE) loss.
+-   Add support for Poisson loss.
 
 ## 1.6.0 2023-09-27
 

tensorflow_decision_forests/keras/keras_test.py

@@ -2233,6 +2233,13 @@ def test_gbt_loss(self):
     )
     mse_model.fit(x=x_train, y=y_train)
 
+    mae_model = keras.GradientBoostedTreesModel(
+        validation_ratio=0.0,
+        loss="MEAN_AVERAGE_ERROR",
+        task=keras.Task.REGRESSION,
+    )
+    mae_model.fit(x=x_train, y=y_train)
+
     multinom_model = keras.GradientBoostedTreesModel(
         validation_ratio=0.0, loss="MULTINOMIAL_LOG_LIKELIHOOD"
     )
@@ -2993,6 +3000,42 @@ def test_monotonic_non_compatible_options(self):
     ):
       model.fit(tf_dataset)
 
+  def test_abalone_mae_loss(self):
+    dataset = abalone_dataset()
+    tf_train, tf_test = dataset_to_tf_dataset(dataset)
+
+    model = keras.GradientBoostedTreesModel(
+        loss="MEAN_AVERAGE_ERROR", task=keras.Task.REGRESSION
+    )
+    model.compile(metrics=["mse", "mae"])
+
+    model.fit(tf_train)
+    evaluation = model.evaluate(tf_test)
+
+    logging.info("Evaluation: %s", evaluation)
+    self.assertLessEqual(evaluation[1], 6.0)
+
+    predictions = model.predict(tf_test)
+    logging.info("Predictions: %s", predictions)
+
+  def test_abalone_poison_loss(self):
+    dataset = abalone_dataset()
+    tf_train, tf_test = dataset_to_tf_dataset(dataset)
+
+    model = keras.GradientBoostedTreesModel(
+        loss="POISSON", task=keras.Task.REGRESSION
+    )
+    model.compile(metrics=["mse", "mae"])
+
+    model.fit(tf_train)
+    evaluation = model.evaluate(tf_test)
+
+    logging.info("Evaluation: %s", evaluation)
+    self.assertLessEqual(evaluation[1], 6.0)
+
+    predictions = model.predict(tf_test)
+    logging.info("Predictions: %s", predictions)
+
 
 if __name__ == "__main__":
   tf.test.main()