add torch.tensor.masked_select

ivy/functional/frontends/torch/tensor.py

@@ -1071,6 +1071,11 @@ def masked_fill_(self, mask, value):
         self.ivy_array = self.masked_fill(mask, value).ivy_array
         return self
 
+    def masked_select(self, mask):
+        x, mask = torch_frontend.broadcast_tensors(self, mask)
+        out = x.reshape(-1).index_select(0, mask.reshape(-1).nonzero().flatten())
+        return torch_frontend.tensor(out)
+
     @with_unsupported_dtypes({"2.2 and below": ("float16", "bfloat16")}, "torch")
     def index_add_(self, dim, index, source, *, alpha=1):
         self.ivy_array = torch_frontend.index_add(
@@ -1891,7 +1896,9 @@ def addcdiv_(self, tensor1, tensor2, *, value=1):
         ).ivy_array
         return self
 
-    @with_unsupported_dtypes({"2.2 and below": ("bfloat16", "float16")}, "torch")
+    @with_supported_dtypes(
+        {"2.2 and below": ("float32", "float64", "complex32", "complex64")}, "torch"
+    )
     def cholesky(self, upper=False):
         return torch_frontend.cholesky(self, upper=upper)
 

ivy_tests/test_ivy/test_frontends/test_torch/test_tensor.py

@@ -9312,6 +9312,41 @@ def test_torch_masked_fill(
     )
 
 
+# masked_select
+@handle_frontend_method(
+    class_tree=CLASS_TREE,
+    init_tree="torch.tensor",
+    method_name="masked_select",
+    x_mask_val=_masked_fill_helper(),
+)
+def test_torch_masked_select(
+    x_mask_val,
+    frontend_method_data,
+    init_flags,
+    method_flags,
+    frontend,
+    on_device,
+    backend_fw,
+):
+    dtype, x, mask, _ = x_mask_val
+    helpers.test_frontend_method(
+        init_input_dtypes=[dtype],
+        backend_to_test=backend_fw,
+        init_all_as_kwargs_np={
+            "data": x,
+        },
+        method_input_dtypes=["bool", dtype],
+        method_all_as_kwargs_np={
+            "mask": mask,
+        },
+        frontend_method_data=frontend_method_data,
+        init_flags=init_flags,
+        method_flags=method_flags,
+        frontend=frontend,
+        on_device=on_device,
+    )
+
+
 # matmul
 @handle_frontend_method(
     class_tree=CLASS_TREE,