def test_einsum():
A = ad.Variable(name="A", shape=[3, 2])
B = ad.Variable(name="B", shape=[2, 3])
y = ad.einsum('ik,kj->ij', A, B)
assert AutodiffParser.parse(y.name, [A, B]).name == y.name
def test_add_3():
A = ad.Variable(name="A", shape=[3])
B = ad.Variable(name="B", shape=[3])
y = A + B + B
assert AutodiffParser.parse(y.name, [A, B]).name == y.name
def sympy_simplify(out, inputs):
"""
Parameters
------------------
out: The to-be simplified node.
inputs: The nodes that are symbols.
------------------
Returns
------------------
out: The simplified output.
------------------
"""
# Retrieve all the leaf nodes to establish parser table.
parser_input = get_all_inputs(out)
cg = CharacterGetter()
input_to_chars = {i: '__' + cg.getchar() for i in inputs}
chars_to_input = dict(zip(input_to_chars.values(), input_to_chars.keys()))
ss_name = f', '.join([input_to_chars[i] for i in inputs])
ss = symbols(ss_name)
formula = out.name
# Visit the inputs in reverse topological order.
all_nodes = reversed(find_topo_sort([out]))
for i in all_nodes:
if i in inputs:
formula = formula.replace(i.name, input_to_chars[i])
ret = str(simplify(formula))
all_nodes = reversed(find_topo_sort([out]))
for i in all_nodes:
if i in inputs:
ret = ret.replace(input_to_chars[i], i.name)
return AutodiffParser.parse(ret, parser_input)
def test_identity():
A = ad.identity(3)
assert AutodiffParser.parse(A.name, []).name == A.name
def test_tensorinv():
A = ad.Variable(name="A", shape=[3, 3])
y = ad.tensorinv(A)
assert AutodiffParser.parse(y.name, [A]).name == y.name
def test_scalar_mul():
A = ad.Variable(name="A", shape=[3])
y = 3.0 * A
assert AutodiffParser.parse(y.name, [A]).name == y.name