Fix Matrix.diff

Author: isuruf

symengine/lib/symengine_wrapper.pyx

@@ -4064,14 +4064,13 @@ import atexit
 atexit.register(module_cleanup)
 
 def diff(expr, *args):
-    cdef Basic ex = sympify(expr)
     cdef Basic prev
     cdef Basic b
     cdef size_t i
     cdef size_t length = len(args)
 
     if not length:
-        return ex
+        return expr
 
     cdef size_t l = 0
     cdef Basic cur_arg, next_arg
@@ -4083,20 +4082,20 @@ def diff(expr, *args):
 
         if l + 1 == length:
             # No next argument, differentiate with no integer argument
-            return ex._diff(cur_arg)
+            return expr._diff(cur_arg)
 
         next_arg = sympify(args[l + 1])
         # Check if the next arg was derivative order
         if isinstance(next_arg, Integer):
             i = int(next_arg)
             for _ in range(i):
-                ex = ex._diff(cur_arg)
+                expr = expr._diff(cur_arg)
             l += 2
             if l == length:
-                return ex
+                return expr
             cur_arg = sympify(args[l])
         else:
-            ex = ex._diff(cur_arg)
+            expr = expr._diff(cur_arg)
             l += 1
             cur_arg = next_arg
 

symengine/tests/test_matrices.py

@@ -655,6 +655,12 @@ def test_cross():
         DenseMatrix(1, 2, [1, 1]).cross(DenseMatrix(1, 2, [1, 1])))
 
 
+def test_diff():
+    x = symbols("x")
+    M = DenseMatrix(1, 2, [x**2, x])
+    assert M.diff(x) == DenseMatrix(1, 2, [2*x, 1])
+
+
 def test_immutablematrix():
     A = ImmutableMatrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
     assert A.shape == (3, 3)