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)(tt.ones((3, 2))).eval() # should work
def __init_group__(self, group):
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 __init__(self,
flow='scale-loc',
local_rv=None,
model=None,
scale_cost_to_minibatch=False,
random_seed=None,
jitter=.1,
**kwargs):
super(NormalizingFlow,
self).__init__(local_rv=local_rv,
scale_cost_to_minibatch=scale_cost_to_minibatch,
model=model,
random_seed=random_seed,
**kwargs)
if isinstance(flow, str):
flow = flows.Formula(flow)(dim=self.global_size,
z0=self.symbolic_initial_global_matrix,
jitter=jitter)
self.gflow = flow
def get_param_spec_for(cls, flow, **kwargs):
return flows.Formula(flow).get_param_spec_for(**kwargs)