Repository: ivy-llc/ivy

Issue:
sequence_mask



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/tensorflow/general_functions.py b/ivy/functional/frontends/tensorflow/general_functions.py
--- a/ivy/functional/frontends/tensorflow/general_functions.py
+++ b/ivy/functional/frontends/tensorflow/general_functions.py
@@ -409,6 +409,34 @@ def searchsorted(sorted_sequence, values, side="left", out_type="int32"):
     return ivy.searchsorted(sorted_sequence, values, side=side, ret_dtype=out_type)
 
 
+@with_supported_dtypes(
+    {"2.14.0 and below": ("int8", "int16", "int32", "int64")}, "tensorflow"
+)
+@to_ivy_arrays_and_back
+def sequence_mask(lengths, maxlen=None, dtype=ivy.bool, name=None):
+    if maxlen is None:
+        maxlen = ivy.maximum(
+            ivy.max(lengths), ivy.max(ivy.arange(ivy.get_num_dims(lengths)))
+        )
+        maxlen = ivy.maximum(0, maxlen)
+    else:
+        maxlen = ivy.array(maxlen)
+    if ivy.get_num_dims(maxlen) is not None and ivy.get_num_dims(maxlen) != 0:
+        raise ValueError(
+            "Argument `maxlen` must be scalar for sequence_mask, "
+            f"received `maxlen` = {maxlen} "
+            f"with shape '{maxlen.get_shape()}' instead"
+        )
+
+    row_vector = ivy.arange(0, int(maxlen), 1)
+    matrix = ivy.expand_dims(lengths, axis=-1)
+    result = row_vector < matrix
+    if dtype is None:
+        return result
+    else:
+        return ivy.astype(result, dtype)
+
+
 @to_ivy_arrays_and_back
 def shape(input, out_type=ivy.int32, name=None):
     out_type = to_ivy_dtype(out_type)


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_tensorflow/test_general_functions.py b/ivy_tests/test_ivy/test_frontends/test_tensorflow/test_general_functions.py
--- a/ivy_tests/test_ivy/test_frontends/test_tensorflow/test_general_functions.py
+++ b/ivy_tests/test_ivy/test_frontends/test_tensorflow/test_general_functions.py
@@ -270,6 +270,40 @@ def _reshape_helper(draw):
     return x, dtype, reshape_shape
 
 
+# sequence_mask
+@st.composite
+def _sequence_mask_helper(draw):
+    max_val = draw(st.integers(min_value=1, max_value=100000))
+    in_dtype, lens = draw(
+        helpers.dtype_and_values(
+            available_dtypes=helpers.get_dtypes("valid"),
+            num_arrays=1,
+            min_value=-max_val,
+            max_value=max_val,
+        )
+    )
+
+    max_len = draw(st.integers(min_value=max_val, max_value=max_val))
+    dtype = draw(
+        st.sampled_from(
+            [
+                "float16",
+                "uint8",
+                "complex128",
+                "bool",
+                "float64",
+                "int8",
+                "int16",
+                "complex64",
+                "float32",
+                "int32",
+                "int64",
+            ]
+        )
+    )
+    return in_dtype, lens, max_len, dtype
+
+
 @st.composite
 def _slice_helper(draw):
     dtype, x, shape = draw(
@@ -1703,6 +1737,34 @@ def test_tensorflow_searchsorted(
     )
 
 
+@handle_frontend_test(
+    fn_tree="tensorflow.sequence_mask",
+    dtype_lens_maxlen=_sequence_mask_helper(),
+    test_with_out=st.just(False),
+)
+def test_tensorflow_sequence_mask(
+    *,
+    dtype_lens_maxlen,
+    frontend,
+    test_flags,
+    fn_tree,
+    on_device,
+    backend_fw,
+):
+    input_dtype, lens, max_len, dtype = dtype_lens_maxlen
+    helpers.test_frontend_function(
+        input_dtypes=input_dtype,
+        backend_to_test=backend_fw,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        lengths=lens[0],
+        maxlen=max_len,
+        dtype=dtype,
+    )
+
+
 # shape
 @handle_frontend_test(
     fn_tree="tensorflow.shape",