guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-02-23 01:57:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix minor issue

Browse source
This commit is contained in:
EiffL 2022-03-26 03:03:59 +01:00
parent 3f9dfa504a
commit 10b093f07f
1 changed files with 3 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
jaxpm/nn.py
View file

@ -35,9 +35,9 @@ class NeuralSplineFourierFilter(hk.Module):
self.n_knots = n_knots self.n_knots = n_knots
self.latent_size = latent_size self.latent_size = latent_size
def __call__(self, k, a): def __call__(self, x, a):
""" """
k: array, scale, normalized to fftfreq default x: array, scale, normalized to fftfreq default
a: scalar, scale factor a: scalar, scale factor
""" """
@ -57,4 +57,4 @@ class NeuralSplineFourierFilter(hk.Module):
# Augment with repeating points # Augment with repeating points
ak = jnp.concatenate([jnp.zeros((3,)), k, jnp.ones((3,))]) ak = jnp.concatenate([jnp.zeros((3,)), k, jnp.ones((3,))])
return _deBoorVectorized(jnp.clip(k/jnp.sqrt(3), 0, 1-1e-4), ak, w, 3) return _deBoorVectorized(jnp.clip(x/jnp.sqrt(3), 0, 1-1e-4), ak, w, 3)