Merge pull request #171 from bjodah/fix-Lambdify-heterogeneous

Fix lambdify heterogeneous

Author: isuruf

symengine/lib/symengine_wrapper.pyx

@@ -3389,25 +3389,24 @@ cdef class _Lambdify(object):
             raise ValueError("Size of out incompatible with number of exprs.")
         self.unsafe_complex(inp, out)
 
-    def __call__(self, inp, out=None):
+    def __call__(self, inp, *, out=None):
         """
         Parameters
         ----------
         inp: array_like
             last dimension must be equal to number of arguments.
         out: array_like or None (default)
-            Allows for for low-overhead use (output argument, must be contiguous).
+            Allows for low-overhead use (output argument, must be contiguous).
             If ``None``: an output container will be allocated (NumPy ndarray).
             If ``len(exprs) > 0`` output is found in the corresponding
-            order. Note that ``out`` is not reshaped.
+            order.
 
         Returns
         -------
         If ``len(exprs) == 1``: ``numpy.ndarray``, otherwise a tuple of such.
 
         """
         cdef:
-            bint reshape_outs
             size_t idx, new_tot_out_size, nbroadcast = 1
             long inp_size
             tuple inp_shape
@@ -3432,12 +3431,10 @@ cdef class _Lambdify(object):
         else:
             inp_shape = inp.shape
         new_tot_out_size = nbroadcast * self.tot_out_size
+        new_out_shapes = [inp_shape[:-1] + out_shape for out_shape in self.out_shapes]
         if out is None:
-            new_out_shapes = [inp_shape[:-1] + out_shape for out_shape in self.out_shapes]
-            reshape_outs = len(new_out_shapes[0]) > 1
             out = np.empty(new_tot_out_size, dtype=self.numpy_dtype)
         else:
-            reshape_outs = False
             if out.size < new_tot_out_size:
                 raise ValueError("Incompatible size of output argument")
             if not (out.flags['C_CONTIGUOUS'] or out.flags['F_CONTIGUOUS']):
@@ -3460,12 +3457,9 @@ cdef class _Lambdify(object):
                 self.unsafe_complex(cmplx_inp, cmplx_out,
                                     idx*self.args_size, idx*self.tot_out_size)
 
-        if reshape_outs:
-            out = out.reshape((nbroadcast, self.tot_out_size))
-            result = [out[:, self.accum_out_sizes[idx]:self.accum_out_sizes[idx+1]].reshape(
-                new_out_shapes[idx]) for idx in range(self.n_exprs)]
-        else:
-            result = [out]
+        out = out.reshape((nbroadcast, self.tot_out_size))
+        result = [out[:, self.accum_out_sizes[idx]:self.accum_out_sizes[idx+1]].reshape(
+            new_out_shapes[idx]) for idx in range(self.n_exprs)]
 
         if self.n_exprs == 1:
             return result[0]
@@ -3569,11 +3563,11 @@ def LambdifyCSE(args, *exprs, cse=None, concatenate=None, **kwargs):
             n_taken += size
         lmb = Lambdify(tuple(args) + cse_symbs, *new_exprs, **kwargs)
         cse_lambda = Lambdify(args, [expr.xreplace(explicit_subs) for expr in cse_exprs], **kwargs)
-        def cb(inp, out=None, **kw):
+        def cb(inp, *, out=None, **kw):
             cse_vals = cse_lambda(inp, **kw)
             print(inp, cse_vals) # DO-NOT-MERGE!
             new_inp = concatenate((inp, cse_vals), axis=-1)
-            return lmb(new_inp, out, **kw)
+            return lmb(new_inp, out=out, **kw)
 
         return cb
     else:

symengine/tests/test_lambdify.py

@@ -140,24 +140,24 @@ def test_numpy_array_out_exceptions():
     lmb = se.Lambdify(args, exprs)
 
     all_right = np.empty(len(exprs))
-    lmb(inp, all_right)
+    lmb(inp, out=all_right)
 
     too_short = np.empty(len(exprs) - 1)
-    raises(ValueError, lambda: (lmb(inp, too_short)))
+    raises(ValueError, lambda: (lmb(inp, out=too_short)))
 
     wrong_dtype = np.empty(len(exprs), dtype=int)
-    raises(ValueError, lambda: (lmb(inp, wrong_dtype)))
+    raises(ValueError, lambda: (lmb(inp, out=wrong_dtype)))
 
     read_only = np.empty(len(exprs))
     read_only.flags['WRITEABLE'] = False
-    raises(ValueError, lambda: (lmb(inp, read_only)))
+    raises(ValueError, lambda: (lmb(inp, out=read_only)))
 
     all_right_broadcast = np.empty((4, len(exprs)))
     inp_bcast = [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]
-    lmb(np.array(inp_bcast), all_right_broadcast)
+    lmb(np.array(inp_bcast), out=all_right_broadcast)
 
     noncontig_broadcast = np.empty((4, len(exprs), 3)).transpose((1, 2, 0))
-    raises(ValueError, lambda: (lmb(inp_bcast, noncontig_broadcast)))
+    raises(ValueError, lambda: (lmb(inp_bcast, out=noncontig_broadcast)))
 
 
 def test_broadcast():
@@ -337,7 +337,7 @@ def test_jacobian():
     lmb = se.Lambdify(args, jac)
     out = np.empty((2, 2))
     inp = X, Y = 7, 11
-    lmb(inp, out)
+    lmb(inp, out=out)
     assert np.allclose(out, [[3 * X**2 * Y, X**3],
                              [Y + 1, X + 1]])
 
@@ -355,7 +355,7 @@ def test_jacobian__broadcast():
     inp1 = 8, 13
     inp2 = 5, 9
     inp = np.array([inp0, inp1, inp2])
-    lmb(inp, out)
+    lmb(inp, out=out)
     for idx, (X, Y) in enumerate([inp0, inp1, inp2]):
         assert np.allclose(out[idx, ...], [[3 * X**2 * Y, X**3],
                                            [Y + 1, X + 1]])
@@ -380,9 +380,8 @@ def test_excessive_out():
     lmb = se.Lambdify([x], [-x])
     inp = np.ones(1)
     out = np.ones(2)
-    out = lmb(inp, out)
+    _ = lmb(inp, out=out[:inp.size])
     assert np.allclose(inp, [1, 1])
-    assert out.shape == (2,)
     assert out[0] == -1
     assert out[1] == 1
 
@@ -558,3 +557,46 @@ def test_lambdify__sympy():
     import sympy as sp
     _sympy_lambdify_heterogeneous_output(se.lambdify, se.DenseMatrix)
     _sympy_lambdify_heterogeneous_output(sp.lambdify, sp.Matrix)
+
+
+def _test_Lambdify_scalar_vector_matrix(Lambdify):
+    if not have_numpy:
+        return
+    args = x, y = se.symbols('x y')
+    vec = se.DenseMatrix([x+y, x*y])
+    jac = vec.jacobian(se.DenseMatrix(args))
+    f = Lambdify(args, x**y, vec, jac)
+    assert f.n_exprs == 3
+    s, v, m = f([2, 3])
+    print(s, v, m)
+    assert s == 2**3
+    assert np.allclose(v, [[2+3], [2*3]])
+    assert np.allclose(m, [
+        [1, 1],
+        [3, 2]
+    ])
+
+    s2, v2, m2 = f([[2, 3], [5, 7]])
+    assert np.allclose(s2, [2**3, 5**7])
+    assert np.allclose(v2, [
+        [[2+3], [2*3]],
+        [[5+7], [5*7]]
+    ])
+    assert np.allclose(m2, [
+        [
+            [1, 1],
+            [3, 2]
+        ],
+        [
+            [1, 1],
+            [7, 5]
+        ]
+    ])
+
+
+def test_Lambdify_scalar_vector_matrix():
+    _test_Lambdify_scalar_vector_matrix(lambda *args: se.Lambdify(*args, backend='lambda'))
+    _test_Lambdify_scalar_vector_matrix(lambda *args: se.LambdifyCSE(*args, backend='lambda'))
+    if se.have_llvm:
+        _test_Lambdify_scalar_vector_matrix(lambda *args: se.Lambdify(*args, backend='llvm'))
+        _test_Lambdify_scalar_vector_matrix(lambda *args: se.LambdifyCSE(*args, backend='llvm'))