mirror of
https://github.com/zebrajr/pytorch.git
synced 2026-01-15 12:15:51 +00:00
5fbaf0eaf8
Summary: This PR changes the compiler to correctly emit in-place operators for augmented assignments (`+=` and friends). - To better match the Python AST structure, add an `AugAssign` tree view and make `Assign` apply only to `=` assignments. - Emit those `AugAssign` exprs in the compiler, dispatching to in-place aten ops for tensors and lowering to simple assignments for scalar types. - In order to preserve (suspect) ONNX export semantics, add a pass to lower the in-place operators to out-of-place operators. Pull Request resolved: https://github.com/pytorch/pytorch/pull/13364 Differential Revision: D12899734 Pulled By: suo fbshipit-source-id: bec83be0062cb0235eb129aed78d6110a9e2c146
72 lines
1.5 KiB
Plaintext
72 lines
1.5 KiB
Plaintext
(def
|
|
(ident fn)
|
|
(decl
|
|
(list
|
|
(param
|
|
(ident x)
|
|
(variable (ident Tensor)))
|
|
(param
|
|
(ident y)
|
|
(variable (ident Tensor)))
|
|
(param
|
|
(ident z)
|
|
(variable (ident Tensor))))
|
|
(option))
|
|
(list
|
|
(assign
|
|
(list (variable (ident q)))
|
|
(None))
|
|
(assign
|
|
(list (variable (ident q)))
|
|
(-
|
|
(+
|
|
(variable (ident x))
|
|
(variable (ident y)))
|
|
(apply
|
|
(.
|
|
(variable (ident z))
|
|
(ident sigmoid))
|
|
(list)
|
|
(list))))
|
|
(expression statement
|
|
(list
|
|
(apply
|
|
(variable (ident print))
|
|
(list (variable (ident q)))
|
|
(list))))
|
|
(assign
|
|
(list (variable (ident w)))
|
|
(unary minus
|
|
(variable (ident z))))
|
|
(if
|
|
(and
|
|
(and
|
|
(not (variable (ident x)))
|
|
(not (variable (ident y))))
|
|
(variable (ident z)))
|
|
(list
|
|
(assign
|
|
(list (variable (ident m)))
|
|
(if
|
|
(not (variable (ident z)))
|
|
(variable (ident x))
|
|
(variable (ident y)))))
|
|
(list))
|
|
(while
|
|
(and
|
|
(<
|
|
(variable (ident x))
|
|
(variable (ident y)))
|
|
(>
|
|
(variable (ident y))
|
|
(variable (ident z))))
|
|
(list
|
|
(assign
|
|
(list (variable (ident q)))
|
|
(variable (ident x)))))
|
|
(assert
|
|
(eq (const 1) (const 1))
|
|
(option (string_literal hello)))
|
|
(return
|
|
(list (variable (ident x))))))
|