提交 95cd6bd5 authored 作者: Matthew Rocklin's avatar Matthew Rocklin

add infer_shape method to Solve Op

上级 d52fc2e5
......@@ -655,6 +655,10 @@ class Solve(Op):
#TODO: use the A_structure to go faster
output_storage[0][0] = scipy.linalg.solve(A, b)
def infer_shape(self, node, shapes):
(Ar, Ac), (Br, Bc) = shapes
return [(Ac, Bc)]
solve = Solve() # general solve
#TODO : SolveTriangular
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论