Jax instance method dot implemented

ivy-llc · Aug 23, 2023 · 68cb6f8 · 68cb6f8
1 parent ca2cb4a
commit 68cb6f8
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 0 deletions.
diff --git a/ivy/functional/frontends/jax/array.py b/ivy/functional/frontends/jax/array.py
@@ -88,6 +88,9 @@ def argmax(
             keepdims=keepdims,
         )
 
+    def dot(self, b, *, precision=None):
+        return jax_frontend.numpy.dot(self, b=b, precision=precision)
+
     def conj(self, /):
         return jax_frontend.numpy.conj(self._ivy_array)
 

diff --git a/ivy_tests/test_ivy/test_frontends/test_jax/test_array.py b/ivy_tests/test_ivy/test_frontends/test_jax/test_array.py
@@ -346,6 +346,40 @@ def test_jax_array_argmax(
         on_device=on_device,
     )
 
+    @handle_frontend_method(
+        class_tree=CLASS_TREE,
+        init_tree="jax.numpy.array",
+        method_name="dot",
+        dtype_and_x=helpers.dtype_and_values(
+            available_dtypes=helpers.get_dtypes("float"),
+            num_arrays=2,
+        ),
+    )
+    def test_jax_array_dot(
+        dtype_and_x,
+        on_device,
+        frontend,
+        frontend_method_data,
+        init_flags,
+        method_flags,
+        backend_fw,
+    ):
+        input_dtype, x = dtype_and_x
+        helpers.test_frontend_method(
+            init_input_dtypes=input_dtype,
+            backend_to_test=backend_fw,
+            init_all_as_kwargs_np={
+                "object": x[0],
+            },
+            method_input_dtypes=input_dtype,
+            method_all_as_kwargs_np={"other": x[1]},
+            frontend=frontend,
+            frontend_method_data=frontend_method_data,
+            init_flags=init_flags,
+            method_flags=method_flags,
+            on_device=on_device,
+        )
+
 
 @handle_frontend_method(
     class_tree=CLASS_TREE,