AD test

Author: cometscome

Project.toml

@@ -1,7 +1,7 @@
 name = "LatticeMatrices"
 uuid = "dd6a91e4-736f-4540-ac85-13822ca7b545"
 authors = ["Yuki Nagai <cometscome@gmail.com>"]
-version = "0.2.6"
+version = "0.2.7"
 
 [deps]
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"

ext/AD/AD.jl

@@ -16,6 +16,7 @@ function kernel_Wiltinger!(i, A, dindexer, ::Val{NC1}, ::Val{NC2}, ::Val{nw}) wh
         for ic = 1:NC1
             X = real(A[ic, jc, indices...])
             Y = imag(A[ic, jc, indices...])
+            #A[jc, ic, indices...] = Complex(0.5 * X, -0.5 * Y)  # ∂/∂A
             A[ic, jc, indices...] = Complex(0.5 * X, -0.5 * Y)  # ∂/∂A
         end
     end
@@ -191,7 +192,7 @@ function Enzyme.EnzymeRules.reverse(::RevConfig,
     dAstruct = A.dval isa Base.RefValue ? A.dval[] : A.dval
     dAstruct === nothing && return (nothing, nothing, nothing)
 
-    # 便利ハンドル
+    # Handy handles
     dAval = dAstruct.A
     dBval = dB.val.A
     N1 = Val(A.val.NC1)
@@ -226,7 +227,7 @@ function Enzyme.EnzymeRules.augmented_primal(
     return AugmentedReturn(nothing, A.dval, nothing)
 end
 
-# 出力Aがconstantな場合
+# When output A is constant
 function Enzyme.EnzymeRules.augmented_primal(
     ::RevConfig,
     ::Const{typeof(traceless_antihermitian!)},
@@ -270,11 +271,11 @@ end
 
 function Enzyme.EnzymeRules.reverse(::RevConfig,
     ::Const{typeof(traceless_antihermitian!)},
-    ::Type{<:Const}, _tape,         # ← 第3引数は戻り値のアクティビティ型（Const{Nothing}）
+    ::Type{<:Const}, _tape,         # Third arg is the return activity type (Const{Nothing})
     A::Annotation{<:LatticeMatrix},
     B::Annotation{<:LatticeMatrix})
 
-    # 上流は「出力Aのshadow」に溜まって返ってくる（dAoutは渡されない）
+    # Upstream accumulates in the "shadow of output A" (dAout is not passed)
     dA = A.dval
     dA = dA isa Base.RefValue ? dA[] : dA
     dA === nothing && return (nothing, nothing)
@@ -291,4 +292,4 @@ function Enzyme.EnzymeRules.reverse(::RevConfig,
     # dB += Π_ah,0(dA)
     JACC.parallel_for(Nsites, kernel_traceless_antihermitian_add!, dB.A, dA.A, NC1, nw, idx)
     return (nothing, nothing)
-end
+end

src/Latticeindices.jl

@@ -133,7 +133,7 @@ end
 @generated function delinearize(
     ::DIndexer{D,dims,strides}, L::Integer, offset::Integer=0
 ) where {D,dims,strides}
-    # 返すブロックを組み立て
+    # Build the block to return
     body = Expr(:block,
         :(Base.@_inline_meta true),
         :(r   = Int(L) - 1),
@@ -142,21 +142,21 @@ end
 
     comps = Vector{Any}(undef, D)
 
-    # d = D..2 まで割り算
+    # Divide for d = D..2
     for d = D:-1:2
         sd = Int(strides[d])
         push!(body.args, :(q, r = Base.divrem(r, $(sd))))
         comps[d] = :(q + off + 1)
     end
-    # d = 1 は余り
+    # Remainder for d = 1
     comps[1] = :(r + off + 1)
 
     push!(body.args, Expr(:tuple, comps...))
     return body
 end
 =#
 
-# ラッパ（ここに @inline を付けるのはOK）
+# Wrapper (OK to add @inline here)
 @inline function delinearize(
     idx::DIndexer{D,dims,strides}, L::Integer, ::Val{nw}
 ) where {D,dims,strides,nw}
@@ -198,4 +198,4 @@ export shiftindices
 
 
 export delinearize
-export linearize
+export linearize