Skip to content

Commit a80f499

Browse files
committed
feat: greedy loop to batch via loop fission
feat: support non-constant extra indices fix: only emit new ops if sure feat: propagate bounds and eliminate noops chore: run fmt feat: check for dynamic_slices feat: unary elementwise working chore: comments feat: clear out unwanted no-ops from loop body feat: elementwise ops are completely supported chore: add a note fix: bad rebase feat: generalize index handling test: elementwise loop fission refactor: support dynamicsliceop feat: sliceinfo generalize feat: lift ops by batch op feat: run batching pass feat: generalize handling for non-reshape case
1 parent b882f6e commit a80f499

File tree

9 files changed

+945
-182
lines changed

9 files changed

+945
-182
lines changed

src/enzyme_ad/jax/Implementations/WhileLoopInfo.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ struct WhileLoopInfo {
3939
std::optional<int64_t> getConstantStart();
4040
std::optional<int64_t> getConstantLimit();
4141

42+
Value getInductionVariable() { return op.getBody().front().getArgument(0); }
43+
4244
int64_t getConstantNumIters();
4345
Value getNumIters(OpBuilder &builder);
4446

0 commit comments

Comments
 (0)