examples.tensor_puzzles.proofs.Puzzles.ScatterAdd

            import JaxLean.Stdlib
import examples.tensor_puzzles.generated.PuzzlesScatterAddset_option maxRecDepth 4096set_option maxHeartbeats 2000000namespace JaxLean.Puzzles.ScatterAddtheorem matches_loop (values : Tensor ℝ [3]) (link : Tensor Int32 [3]) :
    Array.scatter_add values link = Loop.scatter_add values link := byvalues:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link = Loop.scatter_add values link
  funext ⟨i, u⟩values:Tensor ℝ [3]link:Tensor Int32 [3]i:Fin 3u:Index []⊢ Array.scatter_add values link (i, u) = Loop.scatter_add values link (i, u)
  cases uunitvalues:Tensor ℝ [3]link:Tensor Int32 [3]i:Fin 3⊢ Array.scatter_add values link (i, PUnit.unit) = Loop.scatter_add values link (i, PUnit.unit)
  fin_cases iunit.«0»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link ((fun i => i) ⟨0, ⋯⟩, PUnit.unit) =
  Loop.scatter_add values link ((fun i => i) ⟨0, ⋯⟩, PUnit.unit)unit.«1»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link ((fun i => i) ⟨1, ⋯⟩, PUnit.unit) =
  Loop.scatter_add values link ((fun i => i) ⟨1, ⋯⟩, PUnit.unit)unit.«2»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link ((fun i => i) ⟨2, ⋯⟩, PUnit.unit) =
  Loop.scatter_add values link ((fun i => i) ⟨2, ⋯⟩, PUnit.unit) <;>unit.«0»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link ((fun i => i) ⟨0, ⋯⟩, PUnit.unit) =
  Loop.scatter_add values link ((fun i => i) ⟨0, ⋯⟩, PUnit.unit)unit.«1»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link ((fun i => i) ⟨1, ⋯⟩, PUnit.unit) =
  Loop.scatter_add values link ((fun i => i) ⟨1, ⋯⟩, PUnit.unit)unit.«2»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link ((fun i => i) ⟨2, ⋯⟩, PUnit.unit) =
  Loop.scatter_add values link ((fun i => i) ⟨2, ⋯⟩, PUnit.unit)
    simp only [Array.fn__where, Array.scatter_add, Loop.fn__where, Loop.loop_masked_sum,
      Loop.scatter_add, Tensor.map, Tensor.scalar, Tensor.reindex, Tensor.scatterSet_vector,
      Fin.sum_univ_succ]unit.«2»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ (if decide (↑(link (0, ())).toInt = ↑2) = true then values (0, ()) else 0) +
    ((if decide (↑(link (Fin.succ 0, ())).toInt = ↑2) = true then values (Fin.succ 0, ()) else 0) +
      ((if decide (↑(link ((Fin.succ 0).succ, ())).toInt = ↑2) = true then values ((Fin.succ 0).succ, ()) else 0) +
        ∑ x, if decide (↑(link (x.succ.succ.succ, ())).toInt = ↑2) = true then values (x.succ.succ.succ, ()) else 0)) =
  if ⟨2, ⋯⟩ = 2 then
    ((0 + if decide (↑(link (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt = 2) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ()) else 0) +
        if decide (↑(link (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt = 2) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ()) else 0) +
      if decide (↑(link (⟨2 + 1 * ↑0, ⋯⟩, ())).toInt = 2) = true then values (⟨2 + 1 * ↑0, ⋯⟩, ()) else 0
  else
    if ⟨2, ⋯⟩ = 1 then
      ((0 + if decide (↑(link (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt = 1) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ()) else 0) +
          if decide (↑(link (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt = 1) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ()) else 0) +
        if decide (↑(link (⟨2 + 1 * ↑0, ⋯⟩, ())).toInt = 1) = true then values (⟨2 + 1 * ↑0, ⋯⟩, ()) else 0
    else
      if ⟨2, ⋯⟩ = 0 then
        ((0 + if decide (↑(link (⟨0 + 1 * ↑0, ⋯⟩, ())).toInt = 0) = true then values (⟨0 + 1 * ↑0, ⋯⟩, ()) else 0) +
            if decide (↑(link (⟨1 + 1 * ↑0, ⋯⟩, ())).toInt = 0) = true then values (⟨1 + 1 * ↑0, ⋯⟩, ()) else 0) +
          if decide (↑(link (⟨2 + 1 * ↑0, ⋯⟩, ())).toInt = 0) = true then values (⟨2 + 1 * ↑0, ⋯⟩, ()) else 0
      else 0
  all_goals norm_num [Fin.ext_iff, Bool.cond_decide, decide_eq_true_iff]unit.«2»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ (if ↑(link (0, ())).toInt = 2 then values (0, ()) else 0) +
    ((if ↑(link (1, ())).toInt = 2 then values (1, ()) else 0) +
      if ↑(link (2, ())).toInt = 2 then values (2, ()) else 0) =
  ((if ↑(link (0, ())).toInt = 2 then values (0, ()) else 0) +
      if ↑(link (1, ())).toInt = 2 then values (1, ()) else 0) +
    if ↑(link (⟨2, ⋯⟩, ())).toInt = 2 then values (⟨2, ⋯⟩, ()) else 0
  all_goals simp []unit.«2»values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ (if ↑(link (0, ())).toInt = 2 then values (0, ()) else 0) +
    ((if ↑(link (1, ())).toInt = 2 then values (1, ()) else 0) +
      if ↑(link (2, ())).toInt = 2 then values (2, ()) else 0) =
  ((if ↑(link (0, ())).toInt = 2 then values (0, ()) else 0) +
      if ↑(link (1, ())).toInt = 2 then values (1, ()) else 0) +
    if ↑(link (2, ())).toInt = 2 then values (2, ()) else 0
  all_goals ringAll goals completed! 🐙

theorem certificate (values : Tensor ℝ [3]) (link : Tensor Int32 [3]) :
    Jaxpr.Program.eval (.cons values (.cons link .nil)) Array.scatter_add_ir =
      Jaxpr.Program.eval (.cons values (.cons link .nil)) Loop.scatter_add_ir := byvalues:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Jaxpr.Program.eval (Jaxpr.Env.cons values (Jaxpr.Env.cons link Jaxpr.Env.nil)) Array.scatter_add_ir =
  Jaxpr.Program.eval (Jaxpr.Env.cons values (Jaxpr.Env.cons link Jaxpr.Env.nil)) Loop.scatter_add_ir
  rw [Array.scatter_add_translation_correct,values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link =
  Jaxpr.Program.eval (Jaxpr.Env.cons values (Jaxpr.Env.cons link Jaxpr.Env.nil)) Loop.scatter_add_irvalues:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link = Loop.scatter_add values link Loop.scatter_add_translation_correctvalues:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link = Loop.scatter_add values linkvalues:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link = Loop.scatter_add values link]values:Tensor ℝ [3]link:Tensor Int32 [3]⊢ Array.scatter_add values link = Loop.scatter_add values link
  exact matches_loop values linkAll goals completed! 🐙end JaxLean.Puzzles.ScatterAdd