Repository: ivy-llc/ivy

Issue:
block_diag



Discussion:

Resolving patch:
diff --git a/ivy/functional/frontends/torch/miscellaneous_ops.py b/ivy/functional/frontends/torch/miscellaneous_ops.py
--- a/ivy/functional/frontends/torch/miscellaneous_ops.py
+++ b/ivy/functional/frontends/torch/miscellaneous_ops.py
@@ -2,6 +2,7 @@
 from ivy.func_wrapper import with_unsupported_dtypes, with_supported_dtypes
 from ivy.functional.frontends.torch.func_wrapper import to_ivy_arrays_and_back
 from ivy.functional.frontends.torch import promote_types_of_torch_inputs
+import ivy.functional.frontends.torch as torch_frontend
 
 
 @to_ivy_arrays_and_back
@@ -413,6 +414,55 @@ def cov(input, /, *, correction=1, fweights=None, aweights=None):
     return ivy.cov(input, ddof=correction, fweights=fweights, aweights=aweights)
 
 
+# TODO: Add Ivy function for block_diag but only scipy.linalg and \
+# and torch supports block_diag currently
+@to_ivy_arrays_and_back
+def block_diag(*tensors):
+    shapes_list = [ivy.shape(t) for t in tensors]
+    # TODO: Add ivy function to return promoted dtype for multiple tensors at once
+    promoted_dtype = ivy.as_ivy_dtype(tensors[0].dtype)
+    for idx in range(1, len(tensors)):
+        promoted_dtype = torch_frontend.promote_types_torch(
+            tensors[idx - 1].dtype, tensors[idx].dtype
+        )
+
+    inp_tensors = [ivy.asarray(t, dtype=promoted_dtype) for t in tensors]
+    tensors_2d = []
+    result_dim_0, result_dim_1 = 0, 0
+    for idx, t_shape in enumerate(shapes_list):
+        dim_0, dim_1 = 1, 1
+        if len(t_shape) > 2:
+            raise ivy.exceptions.IvyError(
+                "Input tensors must have 2 or fewer dimensions."
+                f"Input {idx} has {len(t_shape)} dimensions"
+            )
+        elif len(t_shape) == 2:
+            dim_0, dim_1 = t_shape
+            tensors_2d.append(inp_tensors[idx])
+        elif len(t_shape) == 1:
+            dim_1 = t_shape[0]
+            tensors_2d.append(ivy.reshape(inp_tensors[idx], shape=(dim_0, dim_1)))
+        else:
+            tensors_2d.append(ivy.reshape(inp_tensors[idx], shape=(dim_0, dim_1)))
+
+        result_dim_0 += dim_0
+        result_dim_1 += dim_1
+        shapes_list[idx] = (dim_0, dim_1)
+
+    ret = ivy.zeros((result_dim_0, result_dim_1), dtype=promoted_dtype)
+    ret_dim_0 = 0
+    ret_dim_1 = 0
+    for idx, t_shape in enumerate(shapes_list):
+        dim_0, dim_1 = t_shape
+        ret[ret_dim_0 : ret_dim_0 + dim_0, ret_dim_1 : ret_dim_1 + dim_1] = (
+            ivy.copy_array(tensors_2d[idx])
+        )
+        ret_dim_0 += dim_0
+        ret_dim_1 += dim_1
+
+    return ret
+
+
 @with_supported_dtypes(
     {"2.0.1 and below": ("complex64", "complex128")},
     "torch",


Test patch:
diff --git a/ivy_tests/test_ivy/test_frontends/test_torch/test_miscellaneous_ops.py b/ivy_tests/test_ivy/test_frontends/test_torch/test_miscellaneous_ops.py
--- a/ivy_tests/test_ivy/test_frontends/test_torch/test_miscellaneous_ops.py
+++ b/ivy_tests/test_ivy/test_frontends/test_torch/test_miscellaneous_ops.py
@@ -1559,6 +1559,43 @@ def test_torch_cov(
     )
 
 
+@handle_frontend_test(
+    fn_tree="torch.block_diag",
+    dtype_and_tensors=helpers.dtype_and_values(
+        available_dtypes=helpers.get_dtypes("valid"),
+        num_arrays=st.integers(min_value=1, max_value=10),
+        min_num_dims=0,
+        max_num_dims=2,
+        allow_inf=True,
+    ),
+    test_with_out=st.just(False),
+)
+def test_torch_block_diag(
+    *,
+    dtype_and_tensors,
+    on_device,
+    fn_tree,
+    frontend,
+    test_flags,
+    backend_fw,
+):
+    dtypes, tensors = dtype_and_tensors
+    if isinstance(dtypes, list):  # If more than one value was generated
+        args = {f"x{i}": np.array(t, dtype=dtypes[i]) for i, t in enumerate(tensors)}
+    else:  # If exactly one value was generated
+        args = {"x0": np.array(tensors, dtype=dtypes)}
+    test_flags.num_positional_args = len(tensors)
+    helpers.test_frontend_function(
+        input_dtypes=dtypes,
+        frontend=frontend,
+        test_flags=test_flags,
+        fn_tree=fn_tree,
+        on_device=on_device,
+        backend_to_test=backend_fw,
+        **args,
+    )
+
+
 # view_as_real
 @handle_frontend_test(
     fn_tree="torch.view_as_real",