def step(self, *args, rng_key=None, **kwargs):
if self.svi_state is None:
if rng_key is None:
rng_key = numpyro.prng_key()
self.svi_state = self.init(rng_key, *args, **kwargs)
try:
self.svi_state, loss = jit(self.update)(self.svi_state, *args,
**kwargs)
except TypeError as e:
if "not a valid JAX type" in str(e):
raise TypeError(
"NumPyro backend requires args, kwargs to be arrays or tuples, "
"dicts of arrays.") from e
else:
raise e
params = jit(super(SVI, self).get_params)(self.svi_state)
get_param_store().update(params)
return loss
def step(self, *args, rng_key=None, **kwargs):
if self.svi_state is None:
if rng_key is None:
rng_key = numpyro.sample('svi.init', dist.PRNGIdentity())
self.svi_state = self.init(rng_key, *args, **kwargs)
try:
self.svi_state, loss = jit(self.update)(self.svi_state, *args,
**kwargs)
except TypeError as e:
if 'not a valid JAX type' in str(e):
raise TypeError(
'NumPyro backend requires args, kwargs to be arrays or tuples, '
'dicts of arrays.')
else:
raise e
params = jit(super(SVI, self).get_params)(self.svi_state)
get_param_store().update(params)
return loss