import JaxLean.Stdlib import examples.tensor_puzzles.generated.PuzzlesDiagset_option maxRecDepth 4096set_option maxHeartbeats 2000000namespace JaxLean.Puzzles.Diagtheorem matches_loop (a : Tensor ℝ [3, 3]) : Array.diag a = Loop.diag a := bya:Tensor ℝ [3, 3]⊢ Array.diag a = Loop.diag a funext ⟨i, u⟩a:Tensor ℝ [3, 3]i:Fin 3u:Index []⊢ Array.diag a (i, u) = Loop.diag a (i, u) cases uunita:Tensor ℝ [3, 3]i:Fin 3⊢ Array.diag a (i, PUnit.unit) = Loop.diag a (i, PUnit.unit) fin_cases iunit.«0»a:Tensor ℝ [3, 3]⊢ Array.diag a ((fun i => i) ⟨0, ⋯⟩, PUnit.unit) = Loop.diag a ((fun i => i) ⟨0, ⋯⟩, PUnit.unit)unit.«1»a:Tensor ℝ [3, 3]⊢ Array.diag a ((fun i => i) ⟨1, ⋯⟩, PUnit.unit) = Loop.diag a ((fun i => i) ⟨1, ⋯⟩, PUnit.unit)unit.«2»a:Tensor ℝ [3, 3]⊢ Array.diag a ((fun i => i) ⟨2, ⋯⟩, PUnit.unit) = Loop.diag a ((fun i => i) ⟨2, ⋯⟩, PUnit.unit) <;>unit.«0»a:Tensor ℝ [3, 3]⊢ Array.diag a ((fun i => i) ⟨0, ⋯⟩, PUnit.unit) = Loop.diag a ((fun i => i) ⟨0, ⋯⟩, PUnit.unit)unit.«1»a:Tensor ℝ [3, 3]⊢ Array.diag a ((fun i => i) ⟨1, ⋯⟩, PUnit.unit) = Loop.diag a ((fun i => i) ⟨1, ⋯⟩, PUnit.unit)unit.«2»a:Tensor ℝ [3, 3]⊢ Array.diag a ((fun i => i) ⟨2, ⋯⟩, PUnit.unit) = Loop.diag a ((fun i => i) ⟨2, ⋯⟩, PUnit.unit) simp only [Array.fn__where, Array.diag, Loop.diag, Tensor.map₂, Tensor.scalar, Tensor.reindex, Tensor.scatterSet_vector, Fin.sum_univ_succ]unit.«2»a:Tensor ℝ [3, 3]⊢ (a (⟨2, ⋯⟩, 0, ()) * if decide (↑2 = ↑↑0) = true then 1 else 0) + ((a (⟨2, ⋯⟩, Fin.succ 0, ()) * if decide (↑2 = ↑↑(Fin.succ 0)) = true then 1 else 0) + ((a (⟨2, ⋯⟩, (Fin.succ 0).succ, ()) * if decide (↑2 = ↑↑(Fin.succ 0).succ) = true then 1 else 0) + ∑ x, a (⟨2, ⋯⟩, x.succ.succ.succ, ()) * if decide (↑2 = ↑↑x.succ.succ.succ) = true then 1 else 0)) = if ⟨2, ⋯⟩ = 2 then a (⟨2 + 1 * ↑0, ⋯⟩, ⟨2 + 1 * ↑0, ⋯⟩, ()) else if ⟨2, ⋯⟩ = 1 then a (⟨1 + 1 * ↑0, ⋯⟩, ⟨1 + 1 * ↑0, ⋯⟩, ()) else if ⟨2, ⋯⟩ = 0 then a (⟨0 + 1 * ↑0, ⋯⟩, ⟨0 + 1 * ↑0, ⋯⟩, ()) else 0 all_goals norm_num [Fin.ext_iff, Bool.cond_decide, decide_eq_true_iff]unit.«2»a:Tensor ℝ [3, 3]⊢ a (⟨2, ⋯⟩, 2, ()) = a (⟨2, ⋯⟩, ⟨2, ⋯⟩, ()) all_goals simp []All goals completed! 🐙 theorem certificate (a : Tensor ℝ [3, 3]) : Jaxpr.Program.eval (.cons a .nil) Array.diag_ir = Jaxpr.Program.eval (.cons a .nil) Loop.diag_ir := bya:Tensor ℝ [3, 3]⊢ Jaxpr.Program.eval (Jaxpr.Env.cons a Jaxpr.Env.nil) Array.diag_ir = Jaxpr.Program.eval (Jaxpr.Env.cons a Jaxpr.Env.nil) Loop.diag_ir rw [Array.diag_translation_correct,a:Tensor ℝ [3, 3]⊢ Array.diag a = Jaxpr.Program.eval (Jaxpr.Env.cons a Jaxpr.Env.nil) Loop.diag_ira:Tensor ℝ [3, 3]⊢ Array.diag a = Loop.diag a Loop.diag_translation_correcta:Tensor ℝ [3, 3]⊢ Array.diag a = Loop.diag aa:Tensor ℝ [3, 3]⊢ Array.diag a = Loop.diag a]a:Tensor ℝ [3, 3]⊢ Array.diag a = Loop.diag a exact matches_loop aAll goals completed! 🐙end JaxLean.Puzzles.Diag