import JaxLean.Stdlib import examples.tensor_puzzles.generated.PuzzlesSequenceMaskset_option maxRecDepth 4096set_option maxHeartbeats 2000000namespace JaxLean.Puzzles.SequenceMasktheorem matches_loop (values : Tensor ℝ [2, 3]) (length : Tensor Int32 [2]) : Array.sequence_mask values length = Loop.sequence_mask values length := byvalues:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length = Loop.sequence_mask values length funext ⟨i, j, u⟩values:Tensor ℝ [2, 3]length:Tensor Int32 [2]i:Fin 2j:Fin 3u:Index []⊢ Array.sequence_mask values length (i, j, u) = Loop.sequence_mask values length (i, j, u) cases uunitvalues:Tensor ℝ [2, 3]length:Tensor Int32 [2]i:Fin 2j:Fin 3⊢ Array.sequence_mask values length (i, j, PUnit.unit) = Loop.sequence_mask values length (i, j, PUnit.unit) fin_cases iunit.«0»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]j:Fin 3⊢ Array.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, j, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, j, PUnit.unit)unit.«1»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]j:Fin 3⊢ Array.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, j, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, j, PUnit.unit) <;>unit.«0»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]j:Fin 3⊢ Array.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, j, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, j, PUnit.unit)unit.«1»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]j:Fin 3⊢ Array.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, j, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, j, PUnit.unit) fin_cases junit.«1».«0»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨0, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨0, ⋯⟩, PUnit.unit)unit.«1».«1»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨1, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨1, ⋯⟩, PUnit.unit)unit.«1».«2»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨2, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨2, ⋯⟩, PUnit.unit) <;>unit.«0».«0»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, (fun i => i) ⟨0, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, (fun i => i) ⟨0, ⋯⟩, PUnit.unit)unit.«0».«1»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, (fun i => i) ⟨1, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, (fun i => i) ⟨1, ⋯⟩, PUnit.unit)unit.«0».«2»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, (fun i => i) ⟨2, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨0, ⋯⟩, (fun i => i) ⟨2, ⋯⟩, PUnit.unit)unit.«1».«0»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨0, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨0, ⋯⟩, PUnit.unit)unit.«1».«1»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨1, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨1, ⋯⟩, PUnit.unit)unit.«1».«2»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨2, ⋯⟩, PUnit.unit) = Loop.sequence_mask values length ((fun i => i) ⟨1, ⋯⟩, (fun i => i) ⟨2, ⋯⟩, PUnit.unit) simp only [Array.fn__where, Array.sequence_mask, Loop.fn__where, Loop.loop_sequence_row, Loop.sequence_mask, Tensor.map, Tensor.scalar, Tensor.reindex, Tensor.scatterSet_vector, Tensor.scatterSet_matrix]unit.«1».«2»values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ (if decide (↑2 < ↑(length (⟨1, ⋯⟩, ())).toInt) = true then values (⟨1, ⋯⟩, ⟨2, ⋯⟩, ()) else 0) = if ⟨1, ⋯⟩ = 1 ∧ ⟨2, ⋯⟩ = 2 then if True then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 2) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (2 + 1 * ↑0), ⋯⟩, ()) else 0 else if 2 = 1 then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 1) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (1 + 1 * ↑0), ⋯⟩, ()) else 0 else if 2 = 0 then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 0) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (0 + 1 * ↑0), ⋯⟩, ()) else 0 else 0 else if ⟨1, ⋯⟩ = 1 ∧ ⟨2, ⋯⟩ = 1 then if 1 = 2 then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 2) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (2 + 1 * ↑0), ⋯⟩, ()) else 0 else if True then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 1) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (1 + 1 * ↑0), ⋯⟩, ()) else 0 else if 1 = 0 then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 0) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (0 + 1 * ↑0), ⋯⟩, ()) else 0 else 0 else if ⟨1, ⋯⟩ = 1 ∧ ⟨2, ⋯⟩ = 0 then if 0 = 2 then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 2) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (2 + 1 * ↑0), ⋯⟩, ()) else 0 else if 0 = 1 then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 1) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (1 + 1 * ↑0), ⋯⟩, ()) else 0 else if True then if decide (↑(length (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt > 0) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (0 + 1 * ↑0), ⋯⟩, ()) else 0 else 0 else if ⟨1, ⋯⟩ = 0 ∧ ⟨2, ⋯⟩ = 2 then if True then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 2) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (2 + 1 * ↑0), ⋯⟩, ()) else 0 else if 2 = 1 then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 1) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (1 + 1 * ↑0), ⋯⟩, ()) else 0 else if 2 = 0 then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 0) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (0 + 1 * ↑0), ⋯⟩, ()) else 0 else 0 else if ⟨1, ⋯⟩ = 0 ∧ ⟨2, ⋯⟩ = 1 then if 1 = 2 then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 2) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (2 + 1 * ↑0), ⋯⟩, ()) else 0 else if True then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 1) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (1 + 1 * ↑0), ⋯⟩, ()) else 0 else if 1 = 0 then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 0) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (0 + 1 * ↑0), ⋯⟩, ()) else 0 else 0 else if ⟨1, ⋯⟩ = 0 ∧ ⟨2, ⋯⟩ = 0 then if 0 = 2 then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 2) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (2 + 1 * ↑0), ⋯⟩, ()) else 0 else if 0 = 1 then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 1) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (1 + 1 * ↑0), ⋯⟩, ()) else 0 else if True then if decide (↑(length (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt > 0) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * (0 + 1 * ↑0), ⋯⟩, ()) else 0 else 0 else 0 all_goals norm_num [Fin.ext_iff, Bool.cond_decide, decide_eq_true_iff]All goals completed! 🐙 theorem certificate (values : Tensor ℝ [2, 3]) (length : Tensor Int32 [2]) : Jaxpr.Program.eval (.cons values (.cons length .nil)) Array.sequence_mask_ir = Jaxpr.Program.eval (.cons values (.cons length .nil)) Loop.sequence_mask_ir := byvalues:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Jaxpr.Program.eval (Jaxpr.Env.cons values (Jaxpr.Env.cons length Jaxpr.Env.nil)) Array.sequence_mask_ir = Jaxpr.Program.eval (Jaxpr.Env.cons values (Jaxpr.Env.cons length Jaxpr.Env.nil)) Loop.sequence_mask_ir rw [Array.sequence_mask_translation_correct,values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length = Jaxpr.Program.eval (Jaxpr.Env.cons values (Jaxpr.Env.cons length Jaxpr.Env.nil)) Loop.sequence_mask_irvalues:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length = Loop.sequence_mask values length Loop.sequence_mask_translation_correctvalues:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length = Loop.sequence_mask values lengthvalues:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length = Loop.sequence_mask values length]values:Tensor ℝ [2, 3]length:Tensor Int32 [2]⊢ Array.sequence_mask values length = Loop.sequence_mask values length exact matches_loop values lengthAll goals completed! 🐙end JaxLean.Puzzles.SequenceMask