def differentiable_loss(self, model, guide, *args, **kwargs):
    with (plate(size=self.num_particles) if (self.num_particles > 1) else contextlib.ExitStack()), enum(first_available_dim=(((- self.max_plate_nesting) - 1) if self.max_plate_nesting else None)):
        guide_tr = trace(guide).get_trace(*args, **kwargs)
        model_tr = trace(replay(model, trace=guide_tr)).get_trace(*args, **kwargs)
    guide_terms = terms_from_trace(guide_tr)
    model_terms = terms_from_trace(model_tr)
    if any(guide_terms['plate_to_step'].values()):
        raise NotImplementedError('TraceMarkovEnum_ELBO does not yet support guide side Markov enumeration')
    with funsor.terms.lazy:
        (contracted_factors, uncontracted_factors) = ([], [])
        for f in model_terms['log_factors']:
            if model_terms['measure_vars'].intersection(f.inputs):
                contracted_factors.append(f)
            else:
                uncontracted_factors.append(f)
        markov_dims = frozenset({plate for (plate, step) in model_terms['plate_to_step'].items() if step})
        contracted_costs = [(model_terms['scale'] * f) for f in funsor.sum_product.dynamic_partial_sum_product(funsor.ops.logaddexp, funsor.ops.add, (model_terms['log_measures'] + contracted_factors), plate_to_step=model_terms['plate_to_step'], eliminate=(model_terms['measure_vars'] | markov_dims))]
        costs = (contracted_costs + uncontracted_factors)
        costs += [(- f) for f in guide_terms['log_factors']]
        plate_vars = (guide_terms['plate_vars'] | model_terms['plate_vars'])
        elbo = to_funsor(0, output=funsor.Real)
        for cost in costs:
            log_prob = funsor.sum_product.sum_product(funsor.ops.logaddexp, funsor.ops.add, guide_terms['log_measures'], plates=plate_vars, eliminate=((plate_vars | guide_terms['measure_vars']) - frozenset(cost.inputs)))
            elbo_term = funsor.Integrate(log_prob, cost, (guide_terms['measure_vars'] & frozenset(cost.inputs)))
            elbo += elbo_term.reduce(funsor.ops.add, (plate_vars & frozenset(cost.inputs)))
    with funsor.interpretations.memoize():
        return (- to_data(apply_optimizer(elbo)))
def differentiable_loss(self, model, guide, *args, **kwargs):
    with (plate(size=self.num_particles) if (self.num_particles > 1) else contextlib.ExitStack()), enum(first_available_dim=(((- self.max_plate_nesting) - 1) if self.max_plate_nesting else None)):
        guide_tr = trace(guide).get_trace(*args, **kwargs)
        model_tr = trace(replay(model, trace=guide_tr)).get_trace(*args, **kwargs)
    guide_terms = terms_from_trace(guide_tr)
    model_terms = terms_from_trace(model_tr)
    if any(guide_terms['plate_to_step'].values()):
        raise NotImplementedError('TraceMarkovEnum_ELBO does not yet support guide side Markov enumeration')
    with funsor.terms.lazy:
        (contracted_factors, uncontracted_factors) = ([], [])
        for f in model_terms['log_factors']:
            if model_terms['measure_vars'].intersection(f.inputs):
                contracted_factors.append(f)
            else:
                uncontracted_factors.append(f)
        markov_dims = frozenset({plate for (plate, step) in model_terms['plate_to_step'].items() if step})
        contracted_costs = [(model_terms['scale'] * f) for f in funsor.sum_product.dynamic_partial_sum_product(funsor.ops.logaddexp, funsor.ops.add, (model_terms['log_measures'] + contracted_factors), plate_to_step=model_terms['plate_to_step'], eliminate=(model_terms['measure_vars'] | markov_dims))]
        costs = (contracted_costs + uncontracted_factors)
        costs += [(- f) for f in guide_terms['log_factors']]
        plate_vars = (guide_terms['plate_vars'] | model_terms['plate_vars'])
        elbo = to_funsor(0, output=funsor.Real)
        for cost in costs:
            log_prob = funsor.sum_product.sum_product(funsor.ops.logaddexp, funsor.ops.add, guide_terms['log_measures'], plates=plate_vars, eliminate=((plate_vars | guide_terms['measure_vars']) - frozenset(cost.inputs)))
            elbo_term = funsor.Integrate(log_prob, cost, (guide_terms['measure_vars'] & frozenset(cost.inputs)))
            elbo += elbo_term.reduce(funsor.ops.add, (plate_vars & frozenset(cost.inputs)))
    with funsor.interpretations.memoize():
        return (- to_data(apply_optimizer(elbo)))