Applying a macro to a variable

General Usage
15

Maybe I’m still misunderstanding what it is you want, but the following gets the gradient of five randomly generated polynomials. Only the first call results in compilation. It’s basically identical to your MWE. Reading the coefficients from a file should be trivial.

using Tullio
using ForwardDiff
using Zygote

# My preference is not to use any macros, such as @Tullio
function generate_polynomial_2D(IJ)
     Poly(coef, x, y) = @tullio poly[i] := coef[j] * x[i]^I[j] * y[i]^J[j] nograd=I nograd=J
     return Poly
end

function test_polynomial_generation_2d(IJ)
	N = 100
    x = rand(N)
    y = rand(N)
    nb_terms = length(IJ[1])
    coef = rand(nb_terms)
	poly = generate_polynomial_2D(IJ)
	loss = (coef, x, y) -> sum(poly(coef, x, y).^2)
	result = Zygote.gradient(loss, coef, x, y)    # <<<< THIS MUST WORK
    return result
end

for _ in 1:5
    terms = rand(1:100)
    @time test_polynomial_generation_2d((rand(0:10,terms),rand(0:10,terms)))
end

  0.047186 seconds (1.73 k allocations: 86.414 KiB, 97.11% compilation time)
  0.001814 seconds (37 allocations: 11.406 KiB)
  0.000715 seconds (37 allocations: 9.406 KiB)
  0.001958 seconds (37 allocations: 11.406 KiB)
  0.000913 seconds (37 allocations: 9.531 KiB)

I don’t know why you are reluctant to use Tullio here. It seems perfect for your usecase. It evaluates multi-threaded, and more importantly performs symbolic differentiation when it can.

You could write your loss function in Tullio too.

function tullio_loss(coef, I, J, x, y)
     @tullio poly[i] := coef[j] * x[i]^I[j] * y[i]^J[j] nograd=I nograd=J
     @tullio loss := poly[i]^2
     return loss
end

for _ in 1:5
    terms = rand(1:100)
    I, J = rand(0:10,terms), rand(0:10,terms)
    x, y = rand(100), rand(100)
    coeff = rand(terms)
    @time Zygote.gradient(tullio_loss, coeff, I, J, x,y)
end

  1.345314 seconds (987.79 k allocations: 49.346 MiB, 18.11% gc time, 99.79% compilation time: 25% of which was recompilation)
  0.000970 seconds (26 allocations: 4.969 KiB)
  0.000611 seconds (26 allocations: 4.812 KiB)
  0.000535 seconds (26 allocations: 4.797 KiB)
  0.001832 seconds (26 allocations: 5.328 KiB)
2 Likes