fix: test_torch_kron

ivy/functional/backends/tensorflow/experimental/linear_algebra.py

@@ -81,7 +81,50 @@ def kron(
     *,
     out: Optional[Union[tf.Tensor, tf.Variable]] = None,
 ) -> Union[tf.Tensor, tf.Variable]:
-    return tf.experimental.numpy.kron(a, b)
+    a_ndims = a.shape.ndims
+    b_ndims = b.shape.ndims
+
+    if a_ndims <= 2 and b_ndims <= 2:
+        return tf.experimental.numpy.kron(a, b)
+
+    a_shape = tf.shape(a)
+    b_shape = tf.shape(b)
+
+    if a_ndims is None or b_ndims is None:
+        raise ValueError("Cannot determine number of dimensions for input arrays.")
+
+    if a_ndims == 0 or b_ndims == 0:
+        return a * b
+
+    if a_ndims < b_ndims:
+        a = tf.reshape(
+            a, tf.concat([tf.ones(b_ndims - a_ndims, dtype=tf.int32), a_shape], 0)
+        )
+        a_shape = tf.shape(a)
+        a_ndims = b_ndims
+    elif b_ndims < a_ndims:
+        b = tf.reshape(
+            b, tf.concat([tf.ones(a_ndims - b_ndims, dtype=tf.int32), b_shape], 0)
+        )
+        b_shape = tf.shape(b)
+
+    ndims = a_ndims
+    if ndims == 0:
+        return a * b
+
+    alphabet = "abcdefghijklmnopqrstuvwxyz"
+    a_axes = alphabet[:ndims]
+    b_axes = alphabet[ndims : 2 * ndims]
+    einsum_str = f"{a_axes},{b_axes}->" + "".join(
+        [sub[item] for sub in zip(a_axes, b_axes) for item in range(2)]
+    )
+    promoted_dtype = tf.experimental.numpy.result_type(a.dtype, b.dtype)
+    a = tf.cast(a, promoted_dtype)
+    b = tf.cast(b, promoted_dtype)
+    res = tf.einsum(einsum_str, a, b)
+
+    final_shape = [a_shape[i] * b_shape[i] for i in range(ndims)]
+    return tf.reshape(res, final_shape)
 
 
 def matrix_exp(

ivy/functional/frontends/torch/miscellaneous_ops.py

@@ -359,6 +359,7 @@ def histc(input, bins=100, min=0, max=0, *, out=None):
 
 
 @to_ivy_arrays_and_back
+@with_supported_dtypes({"2.2 and below": ("float16", "float32", "float64")}, "torch")
 def kron(input, other, *, out=None):
     return ivy.kron(input, other, out=out)
 

ivy_tests/test_ivy/test_frontends/test_torch/test_miscellaneous_ops.py

@@ -1372,7 +1372,7 @@ def test_torch_histc(
 @handle_frontend_test(
     fn_tree="torch.kron",
     dtype_and_x=helpers.dtype_and_values(
-        available_dtypes=helpers.get_dtypes("float"), num_arrays=2
+        available_dtypes=helpers.get_dtypes("valid"), num_arrays=2
     ),
 )
 def test_torch_kron(
@@ -1386,7 +1386,7 @@ def test_torch_kron(
     input_dtypes, x = dtype_and_x
     input, label = x[0], x[1]
     helpers.test_frontend_function(
-        input_dtypes=["float32"],
+        input_dtypes=input_dtypes,
         backend_to_test=backend_fw,
         frontend=frontend,
         test_flags=test_flags,