BatchGaussianRandomFieldModule
. These parameters will often be (effectively) constant in space/time, which means there are initialized with huge (correlation) values. As such, keeping things stable also required using jnp.expm1
to implement (1 - self.phi**2)
in the GaussianRandomField
.
#292