Skip to content

JIT only bare minimum, no explicit conversion to jax arrays

Sebastien Wertz requested to merge less_jit into master
  • Always convert toys to arrays of floats even though they can be integers for poisson toys, to make sure the NLL is only compiled once (it would be compiled twice because of different input array types)
  • By default, only JIT the model NLL and NLL gradient since they are called the most often
  • Don't trace the Hessian when calling Model.compile
  • No explicit conversion of yield and systmatic arrays to jax arrays

The last two might help mitigate memory usage, see #17

Merge request reports