JaxLean.Stdlib.ScatterRules

JaxLean
Stdlib
ScatterRules

import JaxLean.Core.Tensornamespace JaxLean.Tensor

@[simp] theorem scatterSet_same (out : Tensor R s) (p : Index s) (v : R) :
    scatterSet out p v p = v := byR:Types:List ℕout:Tensor R sp:Index sv:R⊢ out.scatterSet p v p = v simp [scatterSet]All goals completed! 🐙

@[simp] theorem scatterSet_other (out : Tensor R s) (p q : Index s) (v : R) (h : q ≠ p) :
    scatterSet out p v q = out q := byR:Types:List ℕout:Tensor R sp:Index sq:Index sv:Rh:q ≠ p⊢ out.scatterSet p v q = out q simp [scatterSet, h]All goals completed! 🐙

@[simp] theorem scatterAdd_same [Add R] (out : Tensor R s) (p : Index s) (v : R) :
    scatterAdd out p v p = out p + v := byR:Types:List ℕinst✝:Add Rout:Tensor R sp:Index sv:R⊢ out.scatterAdd p v p = out p + v simp [scatterAdd]All goals completed! 🐙

@[simp] theorem scatterAdd_other [Add R] (out : Tensor R s) (p q : Index s) (v : R) (h : q ≠ p) :
    scatterAdd out p v q = out q := byR:Types:List ℕinst✝:Add Rout:Tensor R sp:Index sq:Index sv:Rh:q ≠ p⊢ out.scatterAdd p v q = out q simp [scatterAdd, h]All goals completed! 🐙

Coordinate form for vector updates; keeps Index abstract during simplification.

try 'simp' instead of 'simpa'

Note: This linter can be disabled with `set_option linter.unnecessarySimpa false`
@[simp] theorem scatterSet_vector (out : Tensor R [n]) (p i : Fin n) (v : R) :
    scatterSet out (p, ()) v (i, ()) = if i = p then v else out (i, ()) := byR:Typen:ℕout:Tensor R [n]p:Fin ni:Fin nv:R⊢ out.scatterSet (p, ()) v (i, ()) = if i = p then v else out (i, ())
  by_cases h : i = pposR:Typen:ℕout:Tensor R [n]p:Fin ni:Fin nv:Rh:i = p⊢ out.scatterSet (p, ()) v (i, ()) = if i = p then v else out (i, ())negR:Typen:ℕout:Tensor R [n]p:Fin ni:Fin nv:Rh:¬i = p⊢ out.scatterSet (p, ()) v (i, ()) = if i = p then v else out (i, ())
  ·posR:Typen:ℕout:Tensor R [n]p:Fin ni:Fin nv:Rh:i = p⊢ out.scatterSet (p, ()) v (i, ()) = if i = p then v else out (i, ()) subst iposR:Typen:ℕout:Tensor R [n]p:Fin nv:R⊢ out.scatterSet (p, ()) v (p, ()) = if p = p then v else out (p, ())
    try 'simp' instead of 'simpa'

Note: This linter can be disabled with `set_option linter.unnecessarySimpa false`simpa using scatterSet_same out (p, ()) vAll goals completed! 🐙
  ·negR:Typen:ℕout:Tensor R [n]p:Fin ni:Fin nv:Rh:¬i = p⊢ out.scatterSet (p, ()) v (i, ()) = if i = p then v else out (i, ()) have h' : (i, ()) ≠ (p, ()) := fun e => h (congrArg Prod.fst e)negR:Typen:ℕout:Tensor R [n]p:Fin ni:Fin nv:Rh:¬i = ph':(i, ()) ≠ (p, ())⊢ out.scatterSet (p, ()) v (i, ()) = if i = p then v else out (i, ())
    rw [scatterSet_other out (p, ()) (i, ()) v h',negR:Typen:ℕout:Tensor R [n]p:Fin ni:Fin nv:Rh:¬i = ph':(i, ()) ≠ (p, ())⊢ out (i, ()) = if i = p then v else out (i, ())All goals completed! 🐙 if_neg hnegR:Typen:ℕout:Tensor R [n]p:Fin ni:Fin nv:Rh:¬i = ph':(i, ()) ≠ (p, ())⊢ out (i, ()) = out (i, ())All goals completed! 🐙]All goals completed! 🐙

try 'simp' instead of 'simpa'

Note: This linter can be disabled with `set_option linter.unnecessarySimpa false`

@[simp] theorem scatterSet_matrix (out : Tensor R [n, m])
    (p i : Fin n) (q j : Fin m) (v : R) :
    scatterSet out (p, q, ()) v (i, j, ()) =
      if i = p ∧ j = q then v else out (i, j, ()) := byR:Typen:ℕm:ℕout:Tensor R [n, m]p:Fin ni:Fin nq:Fin mj:Fin mv:R⊢ out.scatterSet (p, q, ()) v (i, j, ()) = if i = p ∧ j = q then v else out (i, j, ())
  by_cases h : i = p ∧ j = qposR:Typen:ℕm:ℕout:Tensor R [n, m]p:Fin ni:Fin nq:Fin mj:Fin mv:Rh:i = p ∧ j = q⊢ out.scatterSet (p, q, ()) v (i, j, ()) = if i = p ∧ j = q then v else out (i, j, ())negR:Typen:ℕm:ℕout:Tensor R [n, m]p:Fin ni:Fin nq:Fin mj:Fin mv:Rh:¬(i = p ∧ j = q)⊢ out.scatterSet (p, q, ()) v (i, j, ()) = if i = p ∧ j = q then v else out (i, j, ())
  ·posR:Typen:ℕm:ℕout:Tensor R [n, m]p:Fin ni:Fin nq:Fin mj:Fin mv:Rh:i = p ∧ j = q⊢ out.scatterSet (p, q, ()) v (i, j, ()) = if i = p ∧ j = q then v else out (i, j, ()) rcases h with ⟨rfl, rfl⟩posR:Typen:ℕm:ℕout:Tensor R [n, m]i:Fin nj:Fin mv:R⊢ out.scatterSet (i, j, ()) v (i, j, ()) = if i = i ∧ j = j then v else out (i, j, ())
    try 'simp' instead of 'simpa'

Note: This linter can be disabled with `set_option linter.unnecessarySimpa false`simpa using scatterSet_same out (i, j, ()) vAll goals completed! 🐙
  ·negR:Typen:ℕm:ℕout:Tensor R [n, m]p:Fin ni:Fin nq:Fin mj:Fin mv:Rh:¬(i = p ∧ j = q)⊢ out.scatterSet (p, q, ()) v (i, j, ()) = if i = p ∧ j = q then v else out (i, j, ()) have h' : (i, j, ()) ≠ (p, q, ()) := fun e =>
      h ⟨congrArg Prod.fst e, congrArg (fun x => x.2.1) e⟩negR:Typen:ℕm:ℕout:Tensor R [n, m]p:Fin ni:Fin nq:Fin mj:Fin mv:Rh:¬(i = p ∧ j = q)h':(i, j, ()) ≠ (p, q, ())⊢ out.scatterSet (p, q, ()) v (i, j, ()) = if i = p ∧ j = q then v else out (i, j, ())
    rw [scatterSet_other out (p, q, ()) (i, j, ()) v h',negR:Typen:ℕm:ℕout:Tensor R [n, m]p:Fin ni:Fin nq:Fin mj:Fin mv:Rh:¬(i = p ∧ j = q)h':(i, j, ()) ≠ (p, q, ())⊢ out (i, j, ()) = if i = p ∧ j = q then v else out (i, j, ())All goals completed! 🐙 if_neg hnegR:Typen:ℕm:ℕout:Tensor R [n, m]p:Fin ni:Fin nq:Fin mj:Fin mv:Rh:¬(i = p ∧ j = q)h':(i, j, ()) ≠ (p, q, ())⊢ out (i, j, ()) = out (i, j, ())All goals completed! 🐙]All goals completed! 🐙

Additive collisions are order-independent in this algebraic model.

theorem scatterAdd_comm [AddCommMonoid R] (out : Tensor R s) (p q : Index s) (a b : R) :
    scatterAdd (scatterAdd out p a) q b = scatterAdd (scatterAdd out q b) p a := byR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:R⊢ (out.scatterAdd p a).scatterAdd q b = (out.scatterAdd q b).scatterAdd p a
  funext iR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index s⊢ (out.scatterAdd p a).scatterAdd q b i = (out.scatterAdd q b).scatterAdd p a i
  simp only [scatterAdd]R:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index s⊢ (if i = q then (if i = p then out i + a else out i) + b else if i = p then out i + a else out i) =
  if i = p then (if i = q then out i + b else out i) + a else if i = q then out i + b else out i
  split_ifsposR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index sh✝¹:i = qh✝:i = p⊢ out i + a + b = out i + b + anegR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index sh✝¹:i = qh✝:¬i = p⊢ out i + b = out i + bposR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index sh✝¹:¬i = qh✝:i = p⊢ out i + a = out i + anegR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index sh✝¹:¬i = qh✝:¬i = p⊢ out i = out i <;>posR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index sh✝¹:i = qh✝:i = p⊢ out i + a + b = out i + b + anegR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index sh✝¹:i = qh✝:¬i = p⊢ out i + b = out i + bposR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index sh✝¹:¬i = qh✝:i = p⊢ out i + a = out i + anegR:Types:List ℕinst✝:AddCommMonoid Rout:Tensor R sp:Index sq:Index sa:Rb:Ri:Index sh✝¹:¬i = qh✝:¬i = p⊢ out i = out i simp_all [add_right_comm]All goals completed! 🐙

end JaxLean.Tensor