Add an increment_last_dim() util function for handling the logit layer with dot product.

PiperOrigin-RevId: 368308163

Author: Neural-Link Team

research/carls/util/BUILD

@@ -0,0 +1,43 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Description:
+# Build rules for the utitility component of CARLS.
+
+# Placeholder for internal Python strict & test compatibility macro.
+
+package(
+    default_visibility = ["//research/carls:internal"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+py_library(
+    name = "array_ops",
+    srcs = ["array_ops.py"],
+    srcs_version = "PY3",
+    deps = [
+        # package tensorflow
+    ],
+)
+
+py_test(
+    name = "array_ops_test",
+    srcs = ["array_ops_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        ":array_ops",
+        # package tensorflow
+    ],
+)

research/carls/util/BUILD.oss

@@ -0,0 +1,41 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Description:
+# Build rules for CARLS APIs in Neural Structured Learning.
+
+package(
+    default_visibility = ["//research/carls:internal"],
+    licenses = ["notice"],  # Apache 2.0
+)
+
+py_library(
+    name = "array_ops",
+    srcs = ["array_ops.py"],
+    srcs_version = "PY3",
+    deps = [
+        # package tensorflow
+    ],
+)
+
+py_test(
+    name = "array_ops_test",
+    size = "small",
+    srcs = ["array_ops_test.py"],
+    python_version = "PY3",
+    srcs_version = "PY3",
+    deps = [
+        # package tensorflow
+    ],
+)

research/carls/util/array_ops.py

@@ -0,0 +1,38 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Array related ops."""
+
+import tensorflow as tf
+
+
+def increment_last_dim(input_tensor: tf.Tensor,
+                       default_value: float) -> tf.Tensor:
+  """Grows the size of last dimension of given `input_tensor` by one.
+
+  Examples:
+    - [[1, 2], [3, 4]] -> [[1, 2, 1], [3, 4, 1]] (default_value = 1).
+    - [1, 2, 3] -> [1, 2, 3, 4] (default_value = 4).
+
+  Args:
+    input_tensor: a float tf.Tensor whose last dimension is to be incremented.
+    default_value: a float value denoting the default value for the increased
+      part.
+
+  Returns:
+    A new `tf.Tensor` with increased last dimension size.
+  """
+  input_tensor = tf.dtypes.cast(input_tensor, tf.float32)
+  inc_tensor = tf.ones(tf.shape(input_tensor)[:-1])
+  inc_tensor = tf.expand_dims(inc_tensor, -1) * default_value
+  return tf.concat([input_tensor, inc_tensor], axis=-1)

research/carls/util/array_ops_test.py

@@ -0,0 +1,35 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for neural_structured_learning.research.carls.util.array_ops."""
+
+from research.carls.util import array_ops
+import tensorflow as tf
+
+
+class ArrayOpsTest(tf.test.TestCase):
+
+  def test_increment_last_dim(self):
+    # 1D case
+    input_tensor = tf.constant([2])
+    new_tensor = array_ops.increment_last_dim(input_tensor, 1)
+    self.assertAllClose([2, 1], new_tensor.numpy())
+
+    # 2D case
+    input_tensor = tf.constant([[1, 2], [3, 4]])
+    new_tensor = array_ops.increment_last_dim(input_tensor, 10)
+    self.assertAllClose([[1, 2, 10], [3, 4, 10]], new_tensor.numpy())
+
+
+if __name__ == '__main__':
+  tf.test.main()