@@ -128,11 +128,13 @@ raising).
128128```
129129
130130This IR has a nested loop, but that won't work nicely for GPUs/TPUs. Even for CPUs, XLA
131- often doens't do a great job with loops. By default, we will attempt to raise loops to a
131+ often doens't do a great job with loops. We will attempt to raise loops to a
132132tensor IR.
133133
134134``` @example raising_stablehlo
135- hlo = @code_hlo compute_attractive_force(positions_ra, masses_ra, 2.0f0)
135+ hlo = @code_hlo compile_options=CompileOptions(;
136+ disable_auto_batching_passes=false
137+ ) compute_attractive_force(positions_ra, masses_ra, 2.0f0)
136138@assert !contains(repr(hlo), "stablehlo.while") #hide
137139hlo
138140```
@@ -142,7 +144,9 @@ the values are identical.
142144
143145``` @example raising_stablehlo
144146y_jl = compute_attractive_force(positions, masses, 2.0f0)
145- y_ra = @jit compute_attractive_force(positions_ra, masses_ra, 2.0f0)
147+ y_ra = @jit compile_options=CompileOptions(;
148+ disable_auto_batching_passes=false
149+ ) compute_attractive_force(positions_ra, masses_ra, 2.0f0)
146150maximum(abs, Array(y_ra) .- y_jl)
147151```
148152
0 commit comments