infer_batch_axis.py

@lru_cache()
def infer_batch_axes(model,
                     get_dummy_input=lambda bsize: {'inp': tf.ones([bsize, 3], dtype='int32'),
                                                    'inp_len': tf.constant([3]*bsize, dtype='int32')},
                     check_at=(3, 5, 7),
                     sess=None, **flags):
    """
    This function attempts to figure out batch dimensions by seeing what axes change on different batch sizes.
    It stands as a monument of hatred to a person who thought that time-major axes order in tensorflow RNN is cool.
    :param model: TranslateModel instance to figure out batch dimensions for.
    :param get_dummy_input: a function(batch_size)->(args,kwargs) that creates input batch of given size
        This function should take a single integer and return two outputs: a list and a dictionary.
        args,kwargs = get_dummy_input(batch_size)
        model.encode(*args,**kwargs)

    :param check_at: a list of batch sizes to be used when figuring out batch axis.
    :param sess: tensorflow session to use. Defaults to a new session
    :returns: model.State structure where each element contains a np.array of indices for batch dimensions
    """
    sess = sess or tf.get_default_session() or tf.Session()
    could_be_batch_dim = None

    for batch_size in check_at:
        batch = get_dummy_input(batch_size)
        symbolic_state = model.encode(batch, **flags)
        state = sess.run(symbolic_state)

        state_shapes = nested_map(lambda v: np.array(v.shape), state)

        # only consider dimensions that are of length batch_size
        is_of_right_shape = nested_map(lambda dim_length: dim_length == batch_size, state_shapes)

        if could_be_batch_dim is None:
            could_be_batch_dim = is_of_right_shape

        else:  # only leave True if dimension could be batch dim before AND is of right shape now
            could_be_batch_dim = nested_map(lambda could_before, is_now: could_before & is_now,
                                            could_be_batch_dim, is_of_right_shape)

    # return dimension indices where could_be_batch_dim is True
    return nested_map(lambda could_be: np.where(could_be)[0], could_be_batch_dim)