# Path: weakly_supervised_control/disentanglement/tensorsketch/utils.py
# def count_leading_whitespace(string):
# def shorten(string, num_lines=4):
# def indent(string, spaces=4):
# def pack(x):
# def shapes_to_zeros(*maybe_typed_shapes):
# def elem_isinstance(lst, cls):
# def compute_fan(kernel):
# def compute_out_dims(in_dims, kernel_size, stride,
#                      padding, output_padding,
#                      dilation):
# def assign_moving_average(target, value, momentum):
#     def __init__(self, python_function):
#     def reset(self):
#     def __call__(self, *args, **kwargs):
# def advanced_function(function):
# def reset_tf_function(tf_function):
# class Function(object):


from collections import OrderedDict
from weakly_supervised_control.disentanglement.tensorsketch import utils as tsu
import tensorflow as tf
            m.train(mode)

    def apply(self, fn, filter_fn=None, targets=None):
        # Light wrapper to parse filter_fn and targets args
        if targets is not None:
            assert filter_fn is None, "Cannot use both filter_fn and targets"
            def filter_fn(m): return isinstance(m, targets)

        elif filter_fn is None:
            def filter_fn(m): return True

        self._apply(fn, filter_fn)

    def _apply(self, fn, filter_fn):
        # Apply fn to children first before applying to parent
        # This ensures that parent can override children's decisions
        # Run in chronological reverse order to get reverse topo+chrono apply
        for m in reversed(self._child_modules.values()):

            m._apply(fn, filter_fn)

        if filter_fn(self):
            fn(self)

    def eval(self):
        self.train(False)

    def build(self, *shapes, once=True):
        if once:
            assert not self.built, "{}.built already True".format(self.name)

# Path: weakly_supervised_control/disentanglement/tensorsketch/utils.py
# def count_leading_whitespace(string):
# def shorten(string, num_lines=4):
# def indent(string, spaces=4):
# def pack(x):
# def shapes_to_zeros(*maybe_typed_shapes):
# def elem_isinstance(lst, cls):
# def compute_fan(kernel):
# def compute_out_dims(in_dims, kernel_size, stride,
#                      padding, output_padding,
#                      dilation):
# def assign_moving_average(target, value, momentum):
#     def __init__(self, python_function):
#     def reset(self):
#     def __call__(self, *args, **kwargs):
# def advanced_function(function):
# def reset_tf_function(tf_function):
# class Function(object):


# Path: weakly_supervised_control/disentanglement/tensorsketch/modules/base.py
from collections import OrderedDict
from weakly_supervised_control.disentanglement.tensorsketch import utils as tsu
import tensorflow as tf

            m.train(mode)

    def apply(self, fn, filter_fn=None, targets=None):
        # Light wrapper to parse filter_fn and targets args
        if targets is not None:
            assert filter_fn is None, "Cannot use both filter_fn and targets"
            def filter_fn(m): return isinstance(m, targets)

        elif filter_fn is None:
            def filter_fn(m): return True

        self._apply(fn, filter_fn)

    def _apply(self, fn, filter_fn):
        # Apply fn to children first before applying to parent
        # This ensures that parent can override children's decisions
        # Run in chronological reverse order to get reverse topo+chrono apply
        for m in reversed(self._child_modules.values()):

            m._apply(fn, filter_fn)

        if filter_fn(self):
            fn(self)

    def eval(self):
        self.train(False)

    def build(self, *shapes, once=True):
        if once:
            assert not self.built, "{}.built already True".format(self.name)

        tensors = tsu.shapes_to_zeros(*shapes)