def test_arrayexpr_derivatives1():
res = array_derive(X, X)
assert res == CodegenArrayPermuteDims(CodegenArrayTensorProduct(I, I),
[0, 2, 1, 3])
cg = CodegenArrayTensorProduct(A, X, B)
res = array_derive(cg, X)
assert res == CodegenArrayPermuteDims(
CodegenArrayTensorProduct(I, A, I, B), [0, 4, 2, 3, 1, 5, 6, 7])
cg = CodegenArrayContraction(X, (0, 1))
res = array_derive(cg, X)
assert res == CodegenArrayContraction(CodegenArrayTensorProduct(I, I),
(1, 3))
cg = CodegenArrayDiagonal(X, (0, 1))
res = array_derive(cg, X)
assert res == CodegenArrayDiagonal(CodegenArrayTensorProduct(I, I), (1, 3))
cg = ElementwiseApplyFunction(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(
CodegenArrayDiagonal(
CodegenArrayTensorProduct(ElementwiseApplyFunction(cos, X), I, I),
(0, 3), (1, 5)))
def _(expr: ElementwiseApplyFunction, x: Expr):
assert get_rank(expr) == 2
assert get_rank(x) == 2
fdiff = expr._get_function_fdiff()
dexpr = array_derive(expr.expr, x)
tp = ArrayTensorProduct(ElementwiseApplyFunction(fdiff, expr.expr), dexpr)
td = ArrayDiagonal(tp, (0, 4), (1, 5))
return td
def test_arrayexpr_derivatives1():
res = array_derive(X, X)
assert res == PermuteDims(ArrayTensorProduct(I, I), [0, 2, 1, 3])
cg = ArrayTensorProduct(A, X, B)
res = array_derive(cg, X)
assert res == PermuteDims(
ArrayTensorProduct(I, A, I, B),
[0, 4, 2, 3, 1, 5, 6, 7])
cg = ArrayContraction(X, (0, 1))
res = array_derive(cg, X)
assert res == ArrayContraction(ArrayTensorProduct(I, I), (1, 3))
cg = ArrayDiagonal(X, (0, 1))
res = array_derive(cg, X)
assert res == ArrayDiagonal(ArrayTensorProduct(I, I), (1, 3))
cg = ElementwiseApplyFunction(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
ElementwiseApplyFunction(cos, X),
I,
I
), (0, 3), (1, 5)))
cg = ArrayElementwiseApplyFunc(sin, X)
res = array_derive(cg, X)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
I,
I,
ArrayElementwiseApplyFunc(cos, X)
), (1, 4), (3, 5)))
res = array_derive(A1, A1)
assert res == PermuteDims(
ArrayTensorProduct(Identity(3), Identity(2), Identity(k)),
[0, 2, 4, 1, 3, 5]
)
cg = ArrayElementwiseApplyFunc(sin, A1)
res = array_derive(cg, A1)
assert res.dummy_eq(ArrayDiagonal(
ArrayTensorProduct(
Identity(3), Identity(2), Identity(k),
ArrayElementwiseApplyFunc(cos, A1)
), (1, 6), (3, 7), (5, 8)
))
def _(expr: ArrayElementwiseApplyFunc):
subexpr = _array2matrix(expr.expr)
if isinstance(subexpr, MatrixExpr):
if subexpr.shape != (1, 1):
d = expr.function.bound_symbols[0]
w = Wild("w", exclude=[d])
p = Wild("p", exclude=[d])
m = expr.function.expr.match(w * d**p)
if m is not None:
return m[w] * HadamardPower(subexpr, m[p])
return ElementwiseApplyFunction(expr.function, subexpr)
else:
return ArrayElementwiseApplyFunc(expr.function, subexpr)
def _(expr: ElementwiseApplyFunction):
subexpr, removed = _remove_trivial_dims(expr.expr)
if subexpr.shape == (1, 1):
# TODO: move this to ElementwiseApplyFunction
return expr.function(subexpr), removed + [0, 1]
return ElementwiseApplyFunction(expr.function, subexpr)
def _(expr: ArrayElementwiseApplyFunc):
subexpr = _array2matrix(expr.expr)
if isinstance(subexpr, MatrixExpr):
return ElementwiseApplyFunction(expr.function, subexpr)
else:
return ArrayElementwiseApplyFunc(expr.function, subexpr)
def test_applyfunc_matrix():
double = Lambda(x, x**2)
expr = ElementwiseApplyFunction(double, Xd)
assert isinstance(expr, ElementwiseApplyFunction)
assert expr.doit() == Xd.applyfunc(lambda x: x**2)
assert expr.shape == (3, 3)
assert expr.func(*expr.args) == expr
assert expr[0, 0] == double(Xd[0, 0])
expr = ElementwiseApplyFunction(double, X)
assert isinstance(expr, ElementwiseApplyFunction)
assert isinstance(expr.doit(), ElementwiseApplyFunction)
assert expr == X.applyfunc(double)
assert expr.func(*expr.args) == expr
expr = ElementwiseApplyFunction(exp, X * Y)
assert expr.expr == X * Y
assert expr.function == exp
assert expr == (X * Y).applyfunc(exp)
assert expr.func(*expr.args) == expr
assert isinstance(X * expr, MatMul)
assert (X * expr).shape == (3, 3)
Z = MatrixSymbol("Z", 2, 3)
assert (Z * expr).shape == (2, 3)
expr = ElementwiseApplyFunction(exp, Z.T) * ElementwiseApplyFunction(
exp, Z)
assert expr.shape == (3, 3)
expr = ElementwiseApplyFunction(exp, Z) * ElementwiseApplyFunction(
exp, Z.T)
assert expr.shape == (2, 2)
raises(
ShapeError,
lambda: ElementwiseApplyFunction(exp, Z) * ElementwiseApplyFunction(
exp, Z))
M = Matrix([[x, y], [z, t]])
expr = ElementwiseApplyFunction(sin, M)
assert isinstance(expr, ElementwiseApplyFunction)
assert expr.function == sin
assert expr.expr == M
assert expr.doit() == M.applyfunc(sin)
assert expr.doit() == Matrix([[sin(x), sin(y)], [sin(z), sin(t)]])
assert expr.func(*expr.args) == expr
expr = ElementwiseApplyFunction(double, Xk)
assert expr.doit() == expr
assert expr.subs(k, 2).shape == (2, 2)
assert (expr * expr).shape == (k, k)
M = MatrixSymbol("M", k, t)
expr2 = M.T * expr * M
assert isinstance(expr2, MatMul)
assert expr2.args[1] == expr
assert expr2.shape == (t, t)
expr3 = expr * M
assert expr3.shape == (k, t)
raises(ShapeError, lambda: M * expr)
expr1 = ElementwiseApplyFunction(lambda x: x + 1, Xk)
expr2 = ElementwiseApplyFunction(lambda x: x, Xk)
assert expr1 != expr2
def test_applyfunc_matrix():
double = Lambda(x, x**2)
expr = ElementwiseApplyFunction(double, Xd)
assert isinstance(expr, ElementwiseApplyFunction)
assert expr.doit() == Xd.applyfunc(lambda x: x**2)
assert expr.shape == (3, 3)
expr = ElementwiseApplyFunction(double, X)
assert isinstance(expr, ElementwiseApplyFunction)
assert isinstance(expr.doit(), ElementwiseApplyFunction)
assert expr == X.applyfunc(double)
expr = ElementwiseApplyFunction(exp, X*Y)
assert expr.expr == X*Y
assert expr.function == exp
assert expr == (X*Y).applyfunc(exp)
assert isinstance(X*expr, MatMul)
assert (X*expr).shape == (3, 3)
Z = MatrixSymbol("Z", 2, 3)
assert (Z*expr).shape == (2, 3)
expr = ElementwiseApplyFunction(exp, Z.T)*ElementwiseApplyFunction(exp, Z)
assert expr.shape == (3, 3)
expr = ElementwiseApplyFunction(exp, Z)*ElementwiseApplyFunction(exp, Z.T)
assert expr.shape == (2, 2)
raises(ShapeError, lambda: ElementwiseApplyFunction(exp, Z)*ElementwiseApplyFunction(exp, Z))
M = Matrix([[x, y], [z, t]])
expr = ElementwiseApplyFunction(sin, M)
assert isinstance(expr, ElementwiseApplyFunction)
assert expr.function == sin
assert expr.expr == M
assert expr.doit() == M.applyfunc(sin)
assert expr.doit() == Matrix([[sin(x), sin(y)], [sin(z), sin(t)]])
expr = ElementwiseApplyFunction(double, Xk)
assert expr.doit() == expr
assert expr.subs(k, 2).shape == (2, 2)
assert (expr*expr).shape == (k, k)
M = MatrixSymbol("M", k, t)
expr2 = M.T*expr*M
assert isinstance(expr2, MatMul)
assert expr2.args[1] == expr
assert expr2.shape == (t, t)
expr3 = expr*M
assert expr3.shape == (k, t)
raises(ShapeError, lambda: M*expr)