def test_flow_formula(formula, length, order):
spec = flows.Formula(formula)
flows_list = spec.flows
assert len(flows_list) == length
if order is not None:
assert flows_list == order
spec(dim=2, jitter=1)(at.ones((3, 2))).eval() # should work
def __init_group__(self, group):
raise NotImplementedInference(
"Normalizing flows are not yet ported to v4")
super().__init_group__(group)
# objects to be resolved
# 1. string formula
# 2. not changed default value
# 3. Formula
formula = self._kwargs.get("flow", self._vfam)
jitter = self._kwargs.get("jitter", 1)
if formula is None or isinstance(formula, str):
# case 1 and 2
has_params = self._check_user_params(f=formula)
elif isinstance(formula, flows.Formula):
# case 3
has_params = self._check_user_params(f=formula.formula)
else:
raise TypeError(
"Wrong type provided for NormalizingFlow as `flow` argument, "
"expected Formula or string")
if not has_params:
if formula is None:
formula = self.default_flow
else:
formula = "-".join(
flows.flow_for_params(self.user_params[i]).short_name
for i in range(len(self.user_params)))
if not isinstance(formula, flows.Formula):
formula = flows.Formula(formula)
if self.local:
bs = -1
elif self.batched:
bs = self.bdim
else:
bs = None
self.flow = formula(
dim=self.ddim,
z0=self.symbolic_initial,
jitter=jitter,
params=self.user_params,
batch_size=bs,
)
self._finalize_init()
def get_param_spec_for(cls, flow, **kwargs):
return flows.Formula(flow).get_param_spec_for(**kwargs)