HamletWantToCode

using Flux
# using DifferentialEquations
using Random: seed!

seed!(32);

x0 = rand(2)

model = Chain(Dense(2, 10, σ), Dense(10, 1)) |> f64
f(x) = model(x)

HamletWantToCode

# pyplot()
using Plots
default(grid=false) #src
plot(randn(10))
# closeall();gui()
using Parameters, LTVModelsBase, ValueHistories, DSP, ForwardDiff#, JLD, ReverseDiff#, JacProp
const Diff = ForwardDiff
# const RDiff = ReverseDiff
@with_kw struct LinearSys
    A

HamletWantToCode

using Zygote
using Flux

x0 = rand(2)
dz0 = rand(2)  # test data

model = Chain(Dense(2, 5, relu), Dense(5, 1)) |> f64
f(x) = model(x)

df(x) = gradient(z->f(z)[1], x)[1]

HamletWantToCode

function ising2d(N; BC="PBC")
   J = zeros(N^2, N^2)
   for i in 1:N^2-1
     i₁, i₂ = i%N, i÷N
     if i₁==0
       if BC == "PBC"
         J[1+(i₂-1)*N,i] = J[i,1+(i₂-1)*N] = 1.0
       end
       J[(i₂+1)*N,i] = J[i,(i₂+1)*N] = 1.0
       continue

HamletWantToCode

from jax import vmap

def vectorize(kernel):
    def wraps(*args):
        mv_kernel = vmap(kernel, (None, 0, None), 0)
        mm_kernel = vmap(mv_kernel, (None, None, 0), 1)
        return mm_kernel(*args)
    return wraps