Python keras.backend.stack() Examples

def split_heads_2d(self, ip):
        tensor_shape = K.shape(ip)

        # batch, height, width, channels for axis = -1
        tensor_shape = [tensor_shape[i] for i in range(len(self._shape))]

        batch = tensor_shape[0]
        height = tensor_shape[1]
        width = tensor_shape[2]
        channels = tensor_shape[3]

        # Save the spatial tensor dimensions
        self._batch = batch
        self._height = height
        self._width = width

        ret_shape = K.stack([batch, height, width,  self.num_heads, channels // self.num_heads])
        split = K.reshape(ip, ret_shape)
        transpose_axes = (0, 3, 1, 2, 4)
        split = K.permute_dimensions(split, transpose_axes)

        return split 

def _compute_valid_seed_region(self, seq_length):
        positions = K.arange(seq_length)
        half_block_size = self.block_size // 2
        valid_seed_region = K.switch(
            K.all(
                K.stack(
                    [
                        positions >= half_block_size,
                        positions < seq_length - half_block_size,
                    ],
                    axis=-1,
                ),
                axis=-1,
            ),
            K.ones((seq_length,)),
            K.zeros((seq_length,)),
        )
        return K.expand_dims(K.expand_dims(valid_seed_region, axis=0), axis=-1) 

def _compute_valid_seed_region(self, height, width):
        positions = K.concatenate([
            K.expand_dims(K.tile(K.expand_dims(K.arange(height), axis=1), [1, width]), axis=-1),
            K.expand_dims(K.tile(K.expand_dims(K.arange(width), axis=0), [height, 1]), axis=-1),
        ], axis=-1)
        half_block_size = self.block_size // 2
        valid_seed_region = K.switch(
            K.all(
                K.stack(
                    [
                        positions[:, :, 0] >= half_block_size,
                        positions[:, :, 1] >= half_block_size,
                        positions[:, :, 0] < height - half_block_size,
                        positions[:, :, 1] < width - half_block_size,
                    ],
                    axis=-1,
                ),
                axis=-1,
            ),
            K.ones((height, width)),
            K.zeros((height, width)),
        )
        return K.expand_dims(K.expand_dims(valid_seed_region, axis=0), axis=-1) 

def GenerateMCSamples(inp, layers, K_mc=20):
    if K_mc == 1:
        return apply_layers(inp, layers)
    output_list = []
    for _ in xrange(K_mc):
        output_list += [apply_layers(inp, layers)]  # THIS IS BAD!!! we create new dense layers at every call!!!!
    def pack_out(output_list):
        #output = K.pack(output_list) # K_mc x nb_batch x nb_classes
        output = K.stack(output_list) # K_mc x nb_batch x nb_classes
        return K.permute_dimensions(output, (1, 0, 2)) # nb_batch x K_mc x nb_classes
    def pack_shape(s):
        s = s[0]
        assert len(s) == 2
        return (s[0], K_mc, s[1])
    out = Lambda(pack_out, output_shape=pack_shape)(output_list)
    return out

# evaluation for classification tasks 

def target_shift(proposals, bboxes_gt):
    """Calculates shift needed to transform proposals to match bboxes_gt"""

    proposals = K.cast(proposals, 'float32')
    bboxes_gt = K.cast(bboxes_gt, 'float32')

    proposals_height = proposals[:, 2] - proposals[:, 0]
    proposals_width = proposals[:, 3] - proposals[:, 1]
    proposals_y = proposals[:, 0] + proposals_height * 0.5
    proposals_x = proposals[:, 1] + proposals_width * 0.5

    bboxes_gt_height = bboxes_gt[:, 2] - bboxes_gt[:, 0]
    bboxes_gt_width = bboxes_gt[:, 3] - bboxes_gt[:, 1]
    bboxes_gt_y = bboxes_gt[:, 0] + bboxes_gt_height * 0.5
    bboxes_gt_x = bboxes_gt[:, 1] + bboxes_gt_width * 0.5

    shift = K.stack([(bboxes_gt_y - proposals_y) / proposals_height,
                     (bboxes_gt_x - proposals_x) / proposals_width,
                     K.log(bboxes_gt_height / proposals_height),
                     K.log(bboxes_gt_width / proposals_width)], axis=1)

    return shift 

def shift_anchors(anchors, shifts):
    """Shift anchors"""

    anchors = K.cast(anchors, 'float32')
    shifts = K.cast(shifts, 'float32')

    height = anchors[:, 2] - anchors[:, 0]
    width = anchors[:, 3] - anchors[:, 1]

    center_y = anchors[:, 0] + height / 2
    center_x = anchors[:, 1] + width / 2

    center_y = shifts[:, 0] * height + center_y
    center_x = shifts[:, 1] * width + center_x

    height = K.exp(shifts[:, 2]) * height
    width = K.exp(shifts[:, 3]) * width

    y1 = center_y - height / 2
    y2 = center_y + height / 2

    x1 = center_x - width / 2
    x2 = center_x + width / 2

    return K.stack([y1, x1, y2, x2], axis=1) 

def _compute_valid_seed_region(self, seq_length):
        positions = K.arange(seq_length)
        half_block_size = self.block_size // 2
        valid_seed_region = K.switch(
            K.all(
                K.stack(
                    [
                        positions >= half_block_size,
                        positions < seq_length - half_block_size,
                    ],
                    axis=-1,
                ),
                axis=-1,
            ),
            K.ones((seq_length,)),
            K.zeros((seq_length,)),
        )
        return K.expand_dims(K.expand_dims(valid_seed_region, axis=0), axis=-1) 

def _compute_valid_seed_region(self, height, width):
        positions = K.concatenate([
            K.expand_dims(K.tile(K.expand_dims(K.arange(height), axis=1), [1, width]), axis=-1),
            K.expand_dims(K.tile(K.expand_dims(K.arange(width), axis=0), [height, 1]), axis=-1),
        ], axis=-1)
        half_block_size = self.block_size // 2
        valid_seed_region = K.switch(
            K.all(
                K.stack(
                    [
                        positions[:, :, 0] >= half_block_size,
                        positions[:, :, 1] >= half_block_size,
                        positions[:, :, 0] < height - half_block_size,
                        positions[:, :, 1] < width - half_block_size,
                    ],
                    axis=-1,
                ),
                axis=-1,
            ),
            K.ones((height, width)),
            K.zeros((height, width)),
        )
        return K.expand_dims(K.expand_dims(valid_seed_region, axis=0), axis=-1) 

def sort_tensor_column(X, col_idx):
    # X - the tensor with shape [batch, time, feature_dim, feature_dim]
    # col_idx - the column index with shape[batch, time, r]
    # this function returns a tensor with selected columns by r, i.e. return a tensor with [batch, time, feature_dim, r]
    # notice that the first dimension batch is usually None
    #n_batch = X.get_shape().as_list()[0]
    n_batch = 4
    n_time = X.get_shape().as_list()[1]
    n_dim = X.get_shape().as_list()[2]
    n_rank = col_idx.get_shape().as_list()[-1]
    Xt = tf.transpose(X, [0,1,3,2])
    Xt = tf.reshape(Xt, [n_batch*n_time, n_dim, n_dim])
    col_idx = tf.reshape(col_idx, [n_batch*n_time, n_rank])
    Xt_list = tf.unstack(Xt, axis=0)
    X_sort_list = [tf.gather_nd(Xt_list[t], col_idx[t,:]) for t in range(len(Xt_list))]
    print('X_sort_list[0].shape='+str(X_sort_list[0].shape))
    X_sort = tf.stack(X_sort_list, axis=0)
    X_sort = tf.reshape(X_sort,[n_batch, n_time, n_rank, n_dim])
    X_sort = tf.transpose(X_sort, [0,1,3,2])
    return X_sort 

def tensor_product_local(X,W):
    #input: X in [batch, T, channel]
    #input: W in [T,T]
    #compute X^T * W * X  (* is multiplication)
    n_channels = X.shape[-1]

    A = K.dot(K.permute_dimensions(X, [0,2,1]), W)
    B = K.batch_dot(A,X)
    #return K.reshape(B,[-1,n_channels*n_channels])
    return B





# def _get_tril_batch(D):
#     # convert each element in the batch to lower triangular matrix
#     # D has the size[ batch, dimension, dimension]

#     mat_list = tf.unstack(D, axis=0)
#     fun_tril = Lambda(lambda x: tf.matrix_band_part(x, -1,0))
#     mat_tril_list = [ fun_tril(x) for x in mat_list ]

#     return tf.stack(mat_tril_list, axis=0) 

def tensor_product_local_lowdim(X,W, tril_idx, scaling_mask):
    # input: X in [batch, T, channel]
    # input: w is a 1D vector with size (time, )
    # compute X^T * W * X  (* is multiplication)
    n_channels = X.shape[-1]
    n_batches = X.shape[0]
    
    A = K.dot(K.permute_dimensions(X, [0,2,1]), W)
    B = K.batch_dot(A,X)
    B = B * scaling_mask
    # B_vec = K.reshape(B, [-1, n_channels*n_channels]) # [batch (None), 1, d**2]
 
    #ii = Lambda(lambda x: tf.tile(tf.range(x)[:, tf.newaxis], (1, n_channels)))(n_batches)

    #B_list = Lambda(lambda x: tf.split(x, tf.shape(X)[0], axis=0))(B)

    #B_vec_lowdim_list = [np.sqrt(2)*tf.gather_nd(x, tril_idx) for x in B_list]
    #B_vec_lowdim = K.stack(B_vec_lowdim_list, axis=0)

    B_vec_lowdim = tf.map_fn(lambda x: tf.gather_nd(x, tril_idx), B)
    # print(B_vec_lowdim.shape)
    return B_vec_lowdim*np.sqrt(2) 

def gram_matrix(x, norm_by_channels=False):
        """Calculate gram matrix used in style loss"""
        
        # Assertions on input
        assert K.ndim(x) == 4, 'Input tensor should be a 4d (B, H, W, C) tensor'
        assert K.image_data_format() == 'channels_last', "Please use channels-last format"        
        
        # Permute channels and get resulting shape
        x = K.permute_dimensions(x, (0, 3, 1, 2))
        shape = K.shape(x)
        B, C, H, W = shape[0], shape[1], shape[2], shape[3]
        
        # Reshape x and do batch dot product
        features = K.reshape(x, K.stack([B, C, H*W]))
        gram = K.batch_dot(features, features, axes=2)
        
        # Normalize with channels, height and width
        gram = gram /  K.cast(C * H * W, x.dtype)
        
        return gram
    
    # Prediction functions
    ###################### 

def gram_matrix(x, norm_by_channels=False):
        """Calculate gram matrix used in style loss"""
        
        # Assertions on input
        assert K.ndim(x) == 4, 'Input tensor should be a 4d (B, H, W, C) tensor'
        assert K.image_data_format() == 'channels_last', "Please use channels-last format"        
        
        # Permute channels and get resulting shape
        x = K.permute_dimensions(x, (0, 3, 1, 2))
        shape = K.shape(x)
        B, C, H, W = shape[0], shape[1], shape[2], shape[3]
        
        # Reshape x and do batch dot product
        features = K.reshape(x, K.stack([B, C, H*W]))
        gram = K.batch_dot(features, features, axes=2)
        
        # Normalize with channels, height and width
        gram = gram /  K.cast(C * H * W, x.dtype)
        
        return gram
    
    # Prediction functions
    ###################### 

def parameter_coordinates(shapes):
  num_params = np.sum([np.prod(shape) for shape in shapes])
  output_dims = [shape[-1] for shape in shapes]

  if not output_dims:
    return []

  output_dim = output_dims[0]
  input_dim = num_params / output_dim
  if not all([dim == output_dim for dim in output_dims]):
    raise ValueError("Can't handle different output dimensions")

  return np.stack([
    np.stack([np.arange(input_dim)] * output_dim).flatten() / float(input_dim),
    np.stack([np.arange(output_dim)] * input_dim).T.flatten() / float(output_dim)
  ], axis=-1) 

def rel_to_abs(self, x):
        shape = K.shape(x)
        shape = [shape[i] for i in range(3)]
        B, Nh, L, = shape
        col_pad = K.zeros(K.stack([B, Nh, L, 1]))
        x = K.concatenate([x, col_pad], axis=3)
        flat_x = K.reshape(x, [B, Nh, L * 2 * L])
        flat_pad = K.zeros(K.stack([B, Nh, L - 1]))
        flat_x_padded = K.concatenate([flat_x, flat_pad], axis=2)
        final_x = K.reshape(flat_x_padded, [B, Nh, L + 1, 2 * L - 1])
        final_x = final_x[:, :, :L, L - 1:]
        return final_x 

def combine_heads_2d(self, inputs):
        # [batch, num_heads, height, width, depth_v // num_heads]
        transposed = K.permute_dimensions(inputs, [0, 2, 3, 1, 4])
        # [batch, height, width, num_heads, depth_v // num_heads]
        shape = K.shape(transposed)
        shape = [shape[i] for i in range(5)]

        a, b = shape[-2:]
        ret_shape = K.stack(shape[:-2] + [a * b])
        # [batch, height, width, depth_v]
        return K.reshape(transposed, ret_shape) 

def yolo_eval(yolo_outputs,
              image_shape,
              max_boxes=10,
              score_threshold=.6,
              iou_threshold=.5):
    """Evaluate YOLO model on given input batch and return filtered boxes."""
    box_xy, box_wh, box_confidence, box_class_probs = yolo_outputs
    boxes = yolo_boxes_to_corners(box_xy, box_wh)
    boxes, scores, classes = yolo_filter_boxes(
        boxes, box_confidence, box_class_probs, threshold=score_threshold)

    # Scale boxes back to original image shape.
    height = image_shape[0]
    width = image_shape[1]
    image_dims = K.stack([height, width, height, width])
    image_dims = K.reshape(image_dims, [1, 4])
    boxes = boxes * image_dims

    # TODO: Something must be done about this ugly hack!
    max_boxes_tensor = K.variable(max_boxes, dtype='int32')
    K.get_session().run(tf.variables_initializer([max_boxes_tensor]))
    nms_index = tf.image.non_max_suppression(
        boxes, scores, max_boxes_tensor, iou_threshold=iou_threshold)
    boxes = K.gather(boxes, nms_index)
    scores = K.gather(scores, nms_index)
    classes = K.gather(classes, nms_index)
    return boxes, scores, classes 

def compute_kernel(x, y, kernel='rbf', **kwargs):
    """
        Computes RBF kernel between x and y.
        # Parameters
            x: Tensor
                Tensor with shape [batch_size, z_dim]
            y: Tensor
                Tensor with shape [batch_size, z_dim]
        # Returns
            returns the computed RBF kernel between x and y
    """
    scales = kwargs.get("scales", [])
    if kernel == "rbf":
        x_size = K.shape(x)[0]
        y_size = K.shape(y)[0]
        dim = K.shape(x)[1]
        tiled_x = K.tile(K.reshape(x, K.stack([x_size, 1, dim])), K.stack([1, y_size, 1]))
        tiled_y = K.tile(K.reshape(y, K.stack([1, y_size, dim])), K.stack([x_size, 1, 1]))
        return K.exp(-K.mean(K.square(tiled_x - tiled_y), axis=2) / K.cast(dim, tf.float32))
    elif kernel == 'raphy':
        scales = K.variable(value=np.asarray(scales))
        squared_dist = K.expand_dims(squared_distance(x, y), 0)
        scales = K.expand_dims(K.expand_dims(scales, -1), -1)
        weights = K.eval(K.shape(scales)[0])
        weights = K.variable(value=np.asarray(weights))
        weights = K.expand_dims(K.expand_dims(weights, -1), -1)
        return K.sum(weights * K.exp(-squared_dist / (K.pow(scales, 2))), 0)
    elif kernel == "multi-scale-rbf":
        sigmas = [1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1, 5, 10, 15, 20, 25, 30, 35, 100, 1e3, 1e4, 1e5, 1e6]

        beta = 1. / (2. * (K.expand_dims(sigmas, 1)))
        distances = squared_distance(x, y)
        s = K.dot(beta, K.reshape(distances, (1, -1)))

        return K.reshape(tf.reduce_sum(tf.exp(-s), 0), K.shape(distances)) / len(sigmas) 

def __init__(self, bar_embedder, context_size, 
                 lstm_size, compile_now=False):

        prev_bars = [Input(shape=(None,), name="context_" + str(i)) 
                            for i in range(context_size)]

        embeddings = [bar_embedder(pb) for pb in prev_bars]
        embed_size = bar_embedder.embedding_size                
        
        embeddings_stacked = Lambda(lambda ls: K.stack(ls, axis=1), 
                           output_shape=(context_size, 
                                         embed_size)
                           )(embeddings)
                        
        # encode        
        embeddings_processed = LSTM(lstm_size)(embeddings_stacked)

        self.params = [context_size, lstm_size]
        self.context_size = context_size
        self.encoding_size = lstm_size     

        super().__init__(inputs=prev_bars, 
                         outputs=embeddings_processed,
                         name=repr(self)) # "RhythmEncoder")  
        
        self.bar_embedder = bar_embedder
        
        if compile_now:
            self.compile_default() 

def call(self, inputs):
        features = inputs[0]
        fltr = inputs[1]

        # Convolution
        output = []  # Stores the parallel filters
        for k in range(self.order):
            output_k = features
            for t in range(self.iterations):
                output_k = self.graph_conv_skip([output_k, features, fltr],
                                                self.channels,
                                                'ARMA_skip_{}{}'.format(k, t),
                                                recurrent_k=k if self.share_weights else None,
                                                recurrent_t=t if self.share_weights else None,
                                                activation=self.gcn_activation,
                                                use_bias=self.use_bias,
                                                kernel_initializer=self.kernel_initializer,
                                                bias_initializer=self.bias_initializer,
                                                kernel_regularizer=self.kernel_regularizer,
                                                bias_regularizer=self.bias_regularizer,
                                                kernel_constraint=self.kernel_constraint,
                                                bias_constraint=self.bias_constraint)
            output.append(output_k)

        # Average pooling
        output = K.stack(output, axis=-1)
        output = K.mean(output, axis=-1)
        output = self.activation(output)
        
        return output 

def scale_boxes(boxes, image_shape):
    """ Scales the predicted boxes in order to be drawable on the image"""
    height = image_shape[0]
    width = image_shape[1]
    image_dims = K.stack([height, width, height, width])
    image_dims = K.reshape(image_dims, [1, 4])
    boxes = boxes * image_dims
    return boxes 

def yolo_eval(yolo_outputs,
              image_shape,
              max_boxes=10,
              score_threshold=.6,
              iou_threshold=.5):
    """Evaluate YOLO model on given input batch and return filtered boxes."""
    box_confidence, box_xy, box_wh, box_class_probs = yolo_outputs
    boxes = yolo_boxes_to_corners(box_xy, box_wh)
    boxes, scores, classes = yolo_filter_boxes(
        box_confidence, boxes, box_class_probs, threshold=score_threshold)
    
    # Scale boxes back to original image shape.
    height = image_shape[0]
    width = image_shape[1]
    image_dims = K.stack([height, width, height, width])
    image_dims = K.reshape(image_dims, [1, 4])
    boxes = boxes * image_dims

    # TODO: Something must be done about this ugly hack!
    max_boxes_tensor = K.variable(max_boxes, dtype='int32')
    K.get_session().run(tf.variables_initializer([max_boxes_tensor]))
    nms_index = tf.image.non_max_suppression(
        boxes, scores, max_boxes_tensor, iou_threshold=iou_threshold)
    boxes = K.gather(boxes, nms_index)
    scores = K.gather(scores, nms_index)
    classes = K.gather(classes, nms_index)
    
    return boxes, scores, classes 

def test_keras_unstack_hack():
    y_true_np = np.random.random([1, 3, 2])
    y_true_np[:, :, 0] = 0
    y_true_np[:, :, 1] = 1

    y_true_keras = K.variable(y_true_np)

    y, u = wtte._keras_unstack_hack(y_true_keras)
    y_true_keras_new = K.stack([y, u], axis=-1)

    np.testing.assert_array_equal(K.eval(y_true_keras_new), y_true_np)

# SANITY CHECK: Use pure Weibull data censored at C(ensoring point).
# Should converge to the generating A(alpha) and B(eta) for each timestep 

def build_chimeranet_model(T, F, C, D, n_blstm_units=500, n_blstm_layers=4):
    inputs = Input(shape=(T, F), name='input')
    blstm_top = inputs
    for i in range(1, n_blstm_layers+1):
        blstm_top = Bidirectional(
            LSTM(n_blstm_units, return_sequences=True),
            name='body_blstm_{}'.format(i)
        )(blstm_top)
    body_linear = Dense(F*D, name='body_linear')(blstm_top)
    body = Reshape((T, F, D), name='body')(body_linear)

    embd_activation = Activation('tanh', name='embedding_activation')(body)
    embedding = Lambda(
        lambda x: K.l2_normalize(x, axis=-1), name='embedding'
    )(embd_activation)
    mask_slices = [
        Lambda(
            lambda x: x[:, :, f, :], name='mask_slice_{}'.format(f+1)
        )(body)
        for f in range(F)
    ]
    mask_linears = [
        Dense(
            C, activation='softmax', name='mask_linear_{}'.format(f+1)
        )(mask_slices[f])
        for f in range(F)
    ]
    mask = Lambda(lambda x: K.stack(x, axis=2), name='mask')(mask_linears)

    model = Model(inputs=inputs, outputs=[embedding, mask])
    return model 

def gram_matrix(x, norm_by_channels=False):
    '''
    Returns the Gram matrix of the tensor x.
    '''
    if K.ndim(x) == 3:
        features = K.batch_flatten(K.permute_dimensions(x, (2, 0, 1)))
        shape = K.shape(x)
        C, H, W = shape[0], shape[1], shape[2]
        gram = K.dot(features, K.transpose(features))
    elif K.ndim(x) == 4:
        # Swap from (H, W, C) to (B, C, H, W)
        x = K.permute_dimensions(x, (0, 3, 1, 2))
        shape = K.shape(x)
        B, C, H, W = shape[0], shape[1], shape[2], shape[3]
        # Reshape as a batch of 2D matrices with vectorized channels
        features = K.reshape(x, K.stack([B, C, H*W]))
        # This is a batch of Gram matrices (B, C, C).
        gram = K.batch_dot(features, features, axes=2)
    else:
        raise ValueError('The input tensor should be either a 3d (H, W, C) or 4d (B, H, W, C) tensor.')
    # Normalize the Gram matrix
    if norm_by_channels:
        denominator = C * H * W # Normalization from Johnson
    else:
        denominator = H * W # Normalization from Google
    gram = gram /  K.cast(denominator, x.dtype)

    return gram 

def generalized_dice(y_true, y_pred, exp):
    GD =[]
    print(K.shape(y_pred))
    smooth = 1.0
    for i in range(y_pred.shape[2]):
        y_true_per_label = K.cast(K.not_equal(y_true[:, :, i], 0), K.floatx())  # Change to binary
        y_pred_per_label = y_pred[:, :, i]  # K.cast(K.not_equal(y_pred[:, :, i], 0), K.floatx())  # Change to binary
        weight = K.pow(1/K.sum(y_true_per_label, axis=1), exp)
        intersection = K.sum(y_true_per_label * y_pred_per_label, axis=1)  # (batch size, number of labels)
        union = K.sum(y_true_per_label + y_pred_per_label, axis=1)  # (batch size, number of labels)
        GD.append(weight * (2. * intersection + smooth) / (union + smooth))
    GD_out = K.stack(GD)
    return GD_out 

def call(self, inputs, **kwargs):

        # Give the inputs names
        rpn_match = inputs[0]
        rpn_bbox = inputs[1]
        anchors = inputs[2]

        # Loop through batches
        proposals_batch = []
        for batch_idx in range(self.batch_size):

            # Get top pre_nms_limit anchor indices
            top_k = tf.nn.top_k(rpn_match[batch_idx, :, 1], self.pre_nms_limit, sorted=True, name="top_anchors")

            # Get top pre_nms_limit scores
            top_scores = K.gather(rpn_match[batch_idx, :, 1], top_k.indices)

            # Get top pre_nms_limit anchor shifts
            top_shifts = K.gather(rpn_bbox[batch_idx], top_k.indices)

            # Get top pre_nms_limit anchors
            top_anchors = K.gather(anchors[batch_idx], top_k.indices)

            # Shift the top anchors
            shifted_anchors = gu.shift_anchors(top_anchors, top_shifts)

            # Get indices of proposals after non-max suppression
            proposal_idxs = tf.image.non_max_suppression(shifted_anchors, top_scores, self.post_nms_limit,
                                                         self.nms_threshold)

            # Pad if necessary
            padding = K.maximum(self.post_nms_limit - K.shape(proposal_idxs)[0], 0)
            proposals = tf.pad(K.gather(shifted_anchors, proposal_idxs), [(0, padding), (0, 0)])

            # Append to list
            proposals_batch.append(proposals)

        return K.stack(proposals_batch, axis=0) 

def _ctdet_decode(hm, reg, wh, k=100, output_stride=4):
  hm = K.sigmoid(hm)
  hm = _nms(hm)
  hm_shape = K.shape(hm)
  reg_shape = K.shape(reg)
  wh_shape = K.shape(wh)
  batch, width, cat = hm_shape[0], hm_shape[2], hm_shape[3]

  hm_flat = K.reshape(hm, (batch, -1))
  reg_flat = K.reshape(reg, (reg_shape[0], -1, reg_shape[-1]))
  wh_flat = K.reshape(wh, (wh_shape[0], -1, wh_shape[-1]))

  def _process_sample(args):
    _hm, _reg, _wh = args
    _scores, _inds = tf.math.top_k(_hm, k=k, sorted=True)
    _classes = K.cast(_inds % cat, 'float32')
    _inds = K.cast(_inds / cat, 'int32')
    _xs = K.cast(_inds % width, 'float32')
    _ys = K.cast(K.cast(_inds / width, 'int32'), 'float32')
    _wh = K.gather(_wh, _inds)
    _reg = K.gather(_reg, _inds)

    _xs = _xs + _reg[..., 0]
    _ys = _ys + _reg[..., 1]

    _x1 = _xs - _wh[..., 0] / 2
    _y1 = _ys - _wh[..., 1] / 2
    _x2 = _xs + _wh[..., 0] / 2
    _y2 = _ys + _wh[..., 1] / 2

    # rescale to image coordinates
    _x1 = output_stride * _x1
    _y1 = output_stride * _y1
    _x2 = output_stride * _x2
    _y2 = output_stride * _y2

    _detection = K.stack([_x1, _y1, _x2, _y2, _scores, _classes], -1)
    return _detection

  detections = K.map_fn(_process_sample, [hm_flat, reg_flat, wh_flat], dtype=K.floatx())
  return detections 

def _make_regular_grids(self, batch_size, height, width):
        # making a single regular grid
        x_linspace = K_linspace(-1., 1., width)
        y_linspace = K_linspace(-1., 1., height)
        x_coordinates, y_coordinates = K_meshgrid(x_linspace, y_linspace)
        x_coordinates = K.flatten(x_coordinates)
        y_coordinates = K.flatten(y_coordinates)
        ones = K.ones_like(x_coordinates)
        grid = K.concatenate([x_coordinates, y_coordinates, ones], 0)

        # repeating grids for each batch
        grid = K.flatten(grid)
        grids = K.tile(grid, K.stack([batch_size]))
        return K.reshape(grids, (batch_size, 3, height * width)) 

def _stack_permute_axis_zero(xs):
    xs = Lambda(lambda xs: K.stack(xs, axis=0))(xs)

    # axes (0, 1) are permuted
    perm = [1, 0] + list(range(2, xs.shape.ndims))
    xs = Lambda(lambda xs: K.permute_dimensions(xs, perm))(xs)
    return xs 

def _to_normal2d(output_batch) -> ds.MultivariateNormalTriL:
    """
    :param output_batch: (n_samples, 5)
    :return
    """

    # mean of x and y
    x_mean = Lambda(lambda o: o[:, 0])(output_batch)
    y_mean = Lambda(lambda o: o[:, 1])(output_batch)

    # std of x and y
    # std is must be 0 or positive
    x_std = Lambda(lambda o: K.exp(o[:, 2]))(output_batch)
    y_std = Lambda(lambda o: K.exp(o[:, 3]))(output_batch)

    # correlation coefficient
    # correlation coefficient range is [-1, 1]
    cor = Lambda(lambda o: K.tanh(o[:, 4]))(output_batch)

    loc = Concatenate()([
        Lambda(lambda x_mean: K.expand_dims(x_mean, 1))(x_mean),
        Lambda(lambda y_mean: K.expand_dims(y_mean, 1))(y_mean)
    ])

    x_var = Lambda(lambda x_std: K.square(x_std))(x_std)
    y_var = Lambda(lambda y_std: K.square(y_std))(y_std)
    xy_cor = Multiply()([x_std, y_std, cor])

    cov = Lambda(lambda inputs: K.stack(inputs, axis=0))(
        [x_var, xy_cor, xy_cor, y_var])
    cov = Lambda(lambda cov: K.permute_dimensions(cov, (1, 0)))(cov)
    cov = Reshape((2, 2))(cov)

    scale_tril = Lambda(lambda cov: tf.cholesky(cov))(cov)
    mvn = ds.MultivariateNormalTriL(loc, scale_tril)

    return mvn 

def get_aelstm_autoencoder(ae_encoders, lstm_encoder, ae_decoder):
    y_enc = [None] * len(ae_encoders)
    inp_size = sum([encoder.layers[0].input_shape[1] for encoder in ae_encoders])
    # Input
    x_in = Input(shape=(inp_size,))
    for enc_idx, ae_enc in enumerate(ae_encoders):
        ae_inp_size = ae_encoders[enc_idx].layers[0].input_shape[1]
        x_i = Lambda(
            lambda x: x[:, enc_idx * ae_inp_size:(enc_idx + 1) * ae_inp_size]
        )(x_in)
        y_enc[enc_idx] = ae_encoders[enc_idx](x_i)

    # Ravel AE output for LSTM input
    try:
        y_enc_flat = Lambda(lambda x: KBack.stack(x, axis=1))(y_enc)
    except TypeError:  # If look_back = 1
        y_enc_flat = Lambda(lambda x: KBack.reshape(x, (-1, 1, y_enc[0].shape[1])))(y_enc[0])
    # y_enc_flat = KBack.stack(y_enc, axis=1)
    # Generate embedding

    y = lstm_encoder(y_enc_flat)
    # Generate reconstruction
    x_hat = ae_decoder(y)
    # Autoencoder Model
    autoencoder = Model(input=x_in, output=[x_hat, y])
    return autoencoder 

def call(self, x):
        shape = K.shape(x)
        stack_shape = K.stack([shape[0], shape[1], shape[2], shape[3] // self.nti, self.nti])
        input_reshaped = K.reshape(x, stack_shape)
        mean_per_group = K.mean(input_reshaped, -1)
        return mean_per_group 

def build_model(char_size=27, dim=64, training=True, **kwargs):
  """Build the model."""
  # Inputs
  # Context: (rules, preds, chars,)
  context = L.Input(shape=(None, None, None,), name='context', dtype='int32')
  query = L.Input(shape=(None,), name='query', dtype='int32')

  var_flat = L.Lambda(lambda x: K.reshape(x, K.stack([-1, K.prod(K.shape(x)[1:])])), name='var_flat')
  flat_ctx = var_flat(context)

  # Onehot embedding
  onehot = L.Embedding(char_size, char_size,
                       embeddings_initializer='identity',
                       trainable=False,
                       mask_zero=True,
                       name='onehot')
  embedded_ctx = onehot(flat_ctx) # (?, rules*preds*chars, char_size)
  embedded_q = onehot(query) # (?, chars, char_size)

  # Read query
  _, *states = L.LSTM(dim, return_state=True, name='query_lstm')(embedded_q)
  # Read context
  out, *states = L.LSTM(dim, return_state=True, name='ctx_lstm')(embedded_ctx, initial_state=states)

  # Prediction
  out = L.concatenate([out]+states, name='final_states')
  out = L.Dense(1, activation='sigmoid', name='out')(out)

  model = Model([context, query], out)
  if training:
    model.compile(loss='binary_crossentropy',
                  optimizer='adam',
                  metrics=['acc'])
  return model 

def ScaleBoxes(boxes, imageShape):
    '''Scales the predicted boxes in order to be
    drawable on the image
    
    Arguments:
        boxes  -- all boxes predicted
        imageShape  -- shape to scale the image
    '''
    height = imageShape[0]
    width = imageShape[1]
    imageDims = K.stack([height, width, height, width])
    imageDims = K.reshape(imageDims, [1, 4])
    boxes = boxes * imageDims # Scaling up
    return boxes 

def ScaleBoxes(boxes, imageShape):
    '''Scales the predicted boxes in order to be
    drawable on the image
    
    Arguments:
        boxes  -- all boxes predicted
        imageShape  -- shape to scale the image
    '''
    height = imageShape[0]
    width = imageShape[1]
    imageDims = K.stack([height, width, height, width])
    imageDims = K.reshape(imageDims, [1, 4])
    boxes = boxes * imageDims # Scaling up
    return boxes 

def yolo_eval(yolo_outputs,
              image_shape,
              max_boxes=10,
              score_threshold=.6,
              iou_threshold=.5):
    """Evaluate YOLO model on given input batch and return filtered boxes."""
    box_confidence, box_xy, box_wh, box_class_probs = yolo_outputs
    boxes = yolo_boxes_to_corners(box_xy, box_wh)
    boxes, scores, classes = yolo_filter_boxes(
        box_confidence, boxes, box_class_probs, threshold=score_threshold)
    
    # Scale boxes back to original image shape.
    height = image_shape[0]
    width = image_shape[1]
    image_dims = K.stack([height, width, height, width])
    image_dims = K.reshape(image_dims, [1, 4])
    boxes = boxes * image_dims

    # TODO: Something must be done about this ugly hack!
    max_boxes_tensor = K.variable(max_boxes, dtype='int32')
    K.get_session().run(tf.variables_initializer([max_boxes_tensor]))
    nms_index = tf.image.non_max_suppression(
        boxes, scores, max_boxes_tensor, iou_threshold=iou_threshold)
    boxes = K.gather(boxes, nms_index)
    scores = K.gather(scores, nms_index)
    classes = K.gather(classes, nms_index)
    
    return boxes, scores, classes 

def call(self, inputs, mask=None):
        input_shape = K.shape(inputs)
        if self.mode == self.MODE_ADD:
            batch_size, seq_len, output_dim = input_shape[0], input_shape[1], input_shape[2]
            pos_input = K.tile(K.expand_dims(K.arange(seq_len), axis=0), [batch_size, 1])
        elif self.mode == self.MODE_CONCAT:
            batch_size, seq_len, output_dim = input_shape[0], input_shape[1], self.output_dim
            pos_input = K.tile(K.expand_dims(K.arange(seq_len), axis=0), [batch_size, 1])
        else:
            output_dim = self.output_dim
            pos_input = inputs
        if K.dtype(pos_input) != K.floatx():
            pos_input = K.cast(pos_input, K.floatx())
        evens = K.arange(output_dim // 2) * 2
        odds = K.arange(output_dim // 2) * 2 + 1
        even_embd = K.sin(
            K.dot(
                K.expand_dims(pos_input, -1),
                K.expand_dims(1.0 / K.pow(
                    10000.0,
                    K.cast(evens, K.floatx()) / K.cast(output_dim, K.floatx())
                ), 0)
            )
        )
        odd_embd = K.cos(
            K.dot(
                K.expand_dims(pos_input, -1),
                K.expand_dims(1.0 / K.pow(
                    10000.0, K.cast((odds - 1), K.floatx()) / K.cast(output_dim, K.floatx())
                ), 0)
            )
        )
        embd = K.stack([even_embd, odd_embd], axis=-1)
        output = K.reshape(embd, [-1, K.shape(inputs)[1], output_dim])
        if self.mode == self.MODE_CONCAT:
            output = K.concatenate([inputs, output], axis=-1)
        if self.mode == self.MODE_ADD:
            output += inputs
        return output 

def scale_boxes(boxes, image_shape):
    """ Scales the predicted boxes in order to be drawable on the image"""
    height = image_shape[0]
    width = image_shape[1]
    image_dims = K.stack([height, width, height, width])
    image_dims = K.reshape(image_dims, [1, 4])
    boxes = boxes * image_dims
    return boxes 

def yolo_eval(yolo_outputs,
              image_shape,
              max_boxes=10,
              score_threshold=.6,
              iou_threshold=.5):
    """Evaluate YOLO model on given input batch and return filtered boxes."""
    box_confidence, box_xy, box_wh, box_class_probs = yolo_outputs
    boxes = yolo_boxes_to_corners(box_xy, box_wh)
    boxes, scores, classes = yolo_filter_boxes(
        box_confidence, boxes, box_class_probs, threshold=score_threshold)
    
    # Scale boxes back to original image shape.
    height = image_shape[0]
    width = image_shape[1]
    image_dims = K.stack([height, width, height, width])
    image_dims = K.reshape(image_dims, [1, 4])
    boxes = boxes * image_dims

    # TODO: Something must be done about this ugly hack!
    max_boxes_tensor = K.variable(max_boxes, dtype='int32')
    K.get_session().run(tf.variables_initializer([max_boxes_tensor]))
    nms_index = tf.image.non_max_suppression(
        boxes, scores, max_boxes_tensor, iou_threshold=iou_threshold)
    boxes = K.gather(boxes, nms_index)
    scores = K.gather(scores, nms_index)
    classes = K.gather(classes, nms_index)
    
    return boxes, scores, classes 

def _crop_local(self, reals, fakes, mask_areas):
        """reals, fakes を mask_areasの領域 で切り抜いたデータを得る
        """
        # print("reals, fakes, masks:", reals, fakes, masks)
        # バッチ毎に分割して処理する
        fakes = tf.split(fakes, self.batch_size)
        reals = tf.split(reals, self.batch_size)
        mask_areas = tf.split(mask_areas, self.batch_size)
        real_locals = []
        fake_locals = []
        for real, fake, mask_area in zip(reals, fakes, mask_areas):
            # １次元目のバッチを示す次元を削除
            fake = K.squeeze(fake, 0)
            real = K.squeeze(real, 0)
            mask_area = K.cast(K.squeeze(mask_area, 0), tf.int32)
            top = mask_area[0]
            left = mask_area[1]
            h = mask_area[2] - top
            w = mask_area[3] - left

            # top = util.tfprint(top, prefix="top_debug")
            # left = util.tfprint(left, prefix="left_debug")
            # h = util.tfprint(h, prefix="h_debug")
            # w = util.tfprint(w, prefix="w_debug")

            fake_local = tf.image.crop_to_bounding_box(
                fake, top, left, h, w)
            fake_locals.append(fake_local)

            real_local = tf.image.crop_to_bounding_box(
                real, top, left, h, w)
            real_locals.append(real_local)

        fake_locals = K.stack(fake_locals)
        real_locals = K.stack(real_locals)
        # print("real_locals, fake_locals", real_locals, fake_locals)
        return [real_locals,  fake_locals] 

def head_offsets_loss(gt_offsets, gt_labels, pred_offsets):
    """ヘッドのオフセット回帰の損失関数
    positive（gt_fg > 0）データのみ評価対象とする

    gt_offsets: 正解オフセット
        [N, R, 4]
    gt_labels: 正解データのラベルID
        [N, R]
    pred_offsets: ラベル毎の予測値
        [N, R, n_labels, 4].
    """

    # 正解データのラベルIDに対応するオフセットのみを損失評価対象とする。
    # 論文には以下のようにあるので、正解ラベルのBBoxのみで良さそう。
    # The second task loss, Lloc, is defined over a tuple of true bounding-box
    # regression targets for class u, v = (vx, vy, vw, vh), and a predicted
    # tuple tu = (tux , tuy , tuw, tuh ), again for class u.
    pos_idx = tf.where(gt_labels > 0)
    i = K.cast(pos_idx[:, 0], tf.int32)
    j = K.cast(pos_idx[:, 1], tf.int32)
    k = K.cast(tf.gather_nd(gt_labels, pos_idx), tf.int32)
    pos_pred_idx = K.stack((i, j, k), axis=1)
    pred_offsets = tf.gather_nd(pred_offsets, pos_pred_idx)
    gt_offsets = tf.gather_nd(gt_offsets, pos_idx)

    loss = offsets_loss(gt_offsets, pred_offsets)
    loss = log.tfprint(loss, "head_offsets_loss")
    return loss 

def head_mask_loss(gt_masks, gt_labels, pred_masks):
    """マスクの損失関数

    gt_masks: 正解データ。
        マスクデータをbboxの領域のみ切り抜いてconfig.mask_out_shapeにリサイズしたデータ。
        [N, R, h, w]
        バイナリマスク
    gt_labels: 正解データのラベルID
        [N, R]
    pred_masks: 予測値
        バイナリマスク
        [N, R, n_labels h, w]
    ※h, w は config.mask_out_shape になる。
    """
    # Positiveなラベルが付与されているRoIのみ評価対象とする
    pos_idx = tf.where(gt_labels > 0)
    i = K.cast(pos_idx[:, 0], tf.int32)
    j = K.cast(pos_idx[:, 1], tf.int32)
    k = K.cast(tf.gather_nd(gt_labels, pos_idx), tf.int32)
    # i = log.tfprint(i, "i:head_mask_loss")
    # j = log.tfprint(j, "j:head_mask_loss")
    # k = log.tfprint(k, "k:head_mask_loss")
    pos_pred_idx = K.stack((i, j, k), axis=1)
    # pos_pred_idx = log.tfprint(pos_pred_idx, "pos_pred_idx:head_mask_loss")
    pred_masks = tf.gather_nd(pred_masks, pos_pred_idx)
    gt_masks = tf.gather_nd(gt_masks, pos_idx)

    loss = K.switch(tf.size(gt_masks) > 0,
                    K.binary_crossentropy(gt_masks, pred_masks),
                    tf.constant(0.0))
    loss = K.mean(loss)
    loss = log.tfprint(loss, "head_mask_loss")
    return loss 

def call(self, x):

        # Calculate the pairwise distances between the codewords and the feature vectors
        x_square = K.sum(x ** 2, axis=3, keepdims=True)
        y_square = K.sum(self.V ** 2, axis=2, keepdims=True)
        dists = x_square + y_square - 2 * K.conv2d(x, self.V, strides=(1, 1), padding='valid')
        dists = K.maximum(dists, 0)

        # Quantize the feature vectors
        quantized_features = K.softmax(- dists / (self.sigmas ** 2))

        # Compile the histogram
        if self.spatial_level == 0:
            histogram = K.mean(quantized_features, [1, 2])
        elif self.spatial_level == 1:
            shape = K.shape(quantized_features)
            mid_1 = K.cast(shape[1] / 2, 'int32')
            mid_2 = K.cast(shape[2] / 2, 'int32')
            histogram1 = K.mean(quantized_features[:, :mid_1, :mid_2, :], [1, 2])
            histogram2 = K.mean(quantized_features[:, mid_1:, :mid_2, :], [1, 2])
            histogram3 = K.mean(quantized_features[:, :mid_1, mid_2:, :], [1, 2])
            histogram4 = K.mean(quantized_features[:, mid_1:, mid_2:, :], [1, 2])
            histogram = K.stack([histogram1, histogram2, histogram3, histogram4], 1)
            histogram = K.reshape(histogram, (-1, 4 * self.N_k))
        else:
            # No other spatial level is currently supported (it is trivial to extend the code)
            assert False

        # Simple trick to avoid rescaling issues
        return histogram * self.N_k 

def lambda3(outs):
	block = K.stack(outs)
	bps = []
	for i in range(80):
		bps.append(block[i * 25: (i+1)*25])
	return bps 

def lambda4(bps):
	recomposed = K.stack(bps)
	recomposed = K.reshape(recomposed, (-1,) + (80, 25, 144))
	return recomposed 

def tensor_product(inputs, st_conv_filter_one, conv_len, stride=1, low_dim=False):
    # input - [batch, time, channels]
    local_size=conv_len
    n_frames = inputs.shape[1]
    n_batches = inputs.shape[0]

    x = ZeroPadding1D((local_size//2))(inputs)
    W = Lambda(lambda x: tf.diag(x))(st_conv_filter_one)

    if not low_dim:
        y = [ tensor_product_local(x[:,i:i+local_size,:],W) for i in range(0,n_frames,stride) ]

        outputs =K.stack(y,axis=1) 
        outputs = K.reshape(outputs, [-1,outputs.shape[1],outputs.shape[-2]*outputs.shape[-1] ])
    else:
        n_channels = inputs.get_shape().as_list()[-1]
        tril_idx = np.stack(np.tril_indices(n_channels), axis=0)
        tril_idx2 = np.squeeze(np.split(tril_idx, tril_idx.shape[1], axis=1))


        scaling_mask = np.expand_dims(np.eye(n_channels) / np.sqrt(2), axis=0)


        y = [ tensor_product_local_lowdim(x[:,i:i+local_size,:],W, tril_idx2, scaling_mask ) for i in range(0,n_frames,stride) ]        
        outputs =K.stack(y,axis=1) 

    return outputs 

def tensor_product_with_mean(inputs, st_conv_filter_one, conv_len, feature, stride=1):
    # input - [batch, time, channels]

    if feature not in ['mean','cov','mean_cov']:
        print('[ERROR]: feature for bilinear pooling is not valid')
        sys.exit()

    local_size=conv_len
    n_frames = inputs.shape[1]
    x = ZeroPadding1D(local_size//2)(inputs)

    
    # compute mean
    mu_list = [weighted_average_local(x[:,i:i+local_size,:], st_conv_filter_one) for i in range(0,n_frames)]
    mu = K.stack(mu_list, axis=1)

    
    if feature=='mean':
        return mu[:,::stride, :]

    # compute variance
    x_centered = inputs-mu
    x_centered = ZeroPadding1D(local_size//2)(x_centered)
    W = Lambda(lambda x: tf.diag(x))(st_conv_filter_one)
    sigma_list = [ tensor_product_local(x_centered[:,i:i+local_size,:],W) for i in range(0,n_frames,stride) ]
    sigma =K.stack(sigma_list,axis=1) 
    sigma = K.reshape(sigma, [-1,sigma.shape[1], sigma.shape[-2]*sigma.shape[-1]])

    if feature == 'cov':
        return sigma

    # concatenate mean and covariance
    mu = mu[:,::stride,:]
    outputs = K.concatenate([mu, sigma])


    if feature == 'mean_cov':
        return outputs 

def call(self, X):
        ### here X is the entire mat with [batch, T, D]
        n_frames = X.get_shape().as_list()[1]
        out = X
        for ii in range(self.in_recur):
            out = K.dot(out, self.factorInMats[ii]) + self.factorInBias[ii]
            out = self.act_fun(out)

        ### now we reach the core tensor
        ### do elementwise multiplication, self.core_diag is broadcasted.
        out = K.reshape(K.abs(self.core_diag), [1, 1, -1])*out*out

        ### now we go to the output dim
        for ii in range(self.out_recur):
            
            out = K.dot(out, self.factorOutMats[-ii-1]) + self.factorOutBias[-ii-1]
            out = self.act_fun(out)


        ### now out is [batch, T, out_dim], we do temporal local pooling
        #### zero padding
        out = ZeroPadding1D((self.time_window_size//2))(out)

        W = tf.reshape(self.conv_filter, [1, -1, 1]) # [1, |Nt|, 1]
        out_pool_list = [ K.sum(out[:, i:i+self.time_window_size, :]*W, axis=1)
                         for i in range(0,n_frames, self.stride)  ]

        out_pool = K.stack(out_pool_list,axis=1)

        return out_pool 

def seasonality_model(thetas, backcast_length, forecast_length, is_forecast):
    p = thetas.get_shape().as_list()[-1]
    p1, p2 = (p // 2, p // 2) if p % 2 == 0 else (p // 2, p // 2 + 1)
    t = linear_space(backcast_length, forecast_length, fwd_looking=is_forecast)
    s1 = K.stack([K.cos(2 * np.pi * i * t) for i in range(p1)], axis=0)
    s2 = K.stack([K.sin(2 * np.pi * i * t) for i in range(p2)], axis=0)
    S = K.concatenate([s1, s2], axis=0)
    S = K.cast(S, np.float32)
    return K.dot(thetas, S) 

def trend_model(thetas, backcast_length, forecast_length, is_forecast):
    p = thetas.shape[-1]
    t = linear_space(backcast_length, forecast_length, fwd_looking=is_forecast)
    T = K.transpose(K.stack([t ** i for i in range(p)], axis=0))
    T = K.cast(T, np.float32)
    return K.dot(thetas, K.transpose(T)) 

def forecast_concatenation(thetas, backcast_length, forecast_length, is_forecast):
    p = thetas.shape[-1]
    t = linear_space(backcast_length, forecast_length, fwd_looking=is_forecast)
    T = K.transpose(K.stack([t ** i for i in range(p)], axis=0))
    T = K.cast(T, np.float32)
    return K.dot(thetas, K.transpose(T)) 

def stacklayers(inputs):
    
    """Capa cuya salida consiste de la salida de las capas en inputs concatenadas."""    
    
    return k.stack([inputs[0],inputs[1],inputs[2],inputs[3]],axis = 4) 

def combineOutput(inputs):
    span_start_probabilities, span_end_probabilities = inputs
    return K.stack([span_start_probabilities, span_end_probabilities], axis=-2) 

def time_distributed_dense(x, w, b=None, dropout=None,
                           input_dim=None, units=None, timesteps=None):
    """Apply `y . w + b` for every temporal slice y of x.
    # Arguments
        x: input tensor.
        w: weight matrix.
        b: optional bias vector.
        dropout: wether to apply dropout (same dropout mask
            for every temporal slice of the input).
        input_dim: integer; optional dimensionality of the input.
        units: integer; optional dimensionality of the output.
        timesteps: integer; optional number of timesteps.
    # Returns
        Output tensor.
    """
    if not input_dim:
        input_dim = K.shape(x)[2]
    if not timesteps:
        timesteps = K.shape(x)[1]
    if not units:
        units = K.shape(w)[1]

    if dropout is not None and 0. < dropout < 1.:
        # apply the same dropout pattern at every timestep
        ones = K.ones_like(K.reshape(x[:, 0, :], (-1, input_dim)))
        dropout_matrix = K.dropout(ones, dropout)
        expanded_dropout_matrix = K.repeat(dropout_matrix, timesteps)
        x = K.in_train_phase(x * expanded_dropout_matrix, x)

    # collapse time dimension and batch dimension together
    x = K.reshape(x, (-1, input_dim))
    x = K.dot(x, w)
    if b:
        x += b
    # reshape to 3D tensor
    if K.backend() == 'tensorflow':
        x = K.reshape(x, K.stack([-1, timesteps, units]))
        x.set_shape([None, None, units])
    else:
        x = K.reshape(x, (-1, timesteps, units))
    return x 

def sequential_concat(x):
    return K.stack([x[0], x[1]]) 

def call(self, inputs, training=None):
    if len(inputs) == 3:
        params, trainable_params, x = inputs
        params = self.merge_params(params, trainable_params)
    elif len(inputs) == 2:
        params, x = inputs
    else:
        raise ValueError("Wrong number of inputs")

    offset = 0
    for layer in self.layers:
      layer_params = params[:, offset:offset + layer["num_params"]]
      offset += layer["num_params"]

      if layer["type"] in ["standard-batchnorm", "batch-renorm"]:
        x = K.stack(x, 0)
        self.mean, self.variance = tf.nn.moments(x, [0, 1, 2])

        if training:
          sample_size = K.prod([K.shape(x)[axis] for axis in [0, 1, 2]])
          sample_size = K.cast(sample_size, dtype='float32')
          unbiased_variance = self.variance * sample_size / (sample_size - (1.0 + layer["epsilon"]))

          self.add_update([
            K.moving_average_update(self.moving_means[layer["name"]], self.mean, layer["momentum"]),
            K.moving_average_update(self.moving_vars[layer["name"]], unbiased_variance, layer["momentum"]),
          ], inputs)

      x = [self.evaluate_layer(layer, layer_params[i], x[i], training) for i in range(self.batch_size)]

    output = K.stack(x, 0)
    output._uses_learning_phase = True
    return output 

def scale_boxes(boxes, image_shape):
    """ Scales the predicted boxes in order to be drawable on the image"""
    height = image_shape[0]
    width = image_shape[1]
    image_dims = K.stack([height, width, height, width])
    image_dims = K.reshape(image_dims, [1, 4])
    boxes = boxes * image_dims
    return boxes 

def yolo_eval(yolo_outputs,
              image_shape,
              max_boxes=10,
              score_threshold=.6,
              iou_threshold=.5):
    """Evaluate YOLO model on given input batch and return filtered boxes."""
    box_confidence, box_xy, box_wh, box_class_probs = yolo_outputs
    boxes = yolo_boxes_to_corners(box_xy, box_wh)
    boxes, scores, classes = yolo_filter_boxes(
        box_confidence, boxes, box_class_probs, threshold=score_threshold)
    
    # Scale boxes back to original image shape.
    height = image_shape[0]
    width = image_shape[1]
    image_dims = K.stack([height, width, height, width])
    image_dims = K.reshape(image_dims, [1, 4])
    boxes = boxes * image_dims

    # TODO: Something must be done about this ugly hack!
    max_boxes_tensor = K.variable(max_boxes, dtype='int32')
    K.get_session().run(tf.variables_initializer([max_boxes_tensor]))
    nms_index = tf.image.non_max_suppression(
        boxes, scores, max_boxes_tensor, iou_threshold=iou_threshold)
    boxes = K.gather(boxes, nms_index)
    scores = K.gather(scores, nms_index)
    classes = K.gather(classes, nms_index)
    
    return boxes, scores, classes 

def call(self, inputs):
        span_begin_probabilities, span_end_probabilities = inputs
        return K.stack([span_begin_probabilities, span_end_probabilities], axis = 1) 

def _time_distributed_dense(x, w, b=None, dropout=None,
                            input_dim=None, output_dim=None,
                            timesteps=None, training=None):
    """Apply `y . w + b` for every temporal slice y of x.

    # Arguments
        x: input tensor.
        w: weight matrix.
        b: optional bias vector.
        dropout: wether to apply dropout (same dropout mask
            for every temporal slice of the input).
        input_dim: integer; optional dimensionality of the input.
        output_dim: integer; optional dimensionality of the output.
        timesteps: integer; optional number of timesteps.
        training: training phase tensor or boolean.

    # Returns
        Output tensor.
    """
    if not input_dim:
        input_dim = K.shape(x)[2]
    if not timesteps:
        timesteps = K.shape(x)[1]
    if not output_dim:
        output_dim = K.int_shape(w)[1]

    if dropout is not None and 0. < dropout < 1.:
        # apply the same dropout pattern at every timestep
        ones = K.ones_like(K.reshape(x[:, 0, :], (-1, input_dim)))
        dropout_matrix = K.dropout(ones, dropout)
        expanded_dropout_matrix = K.repeat(dropout_matrix, timesteps)
        x = K.in_train_phase(x * expanded_dropout_matrix, x, training=training)

    # collapse time dimension and batch dimension together
    x = K.reshape(x, (-1, input_dim))
    x = K.dot(x, w)
    if b is not None:
        x = K.bias_add(x, b)
    # reshape to 3D tensor
    if K.backend() == 'tensorflow':
        x = K.reshape(x, K.stack([-1, timesteps, output_dim]))
        x.set_shape([None, None, output_dim])
    else:
        x = K.reshape(x, (-1, timesteps, output_dim))
    return x 

def _time_distributed_dense(x, w, b=None, dropout=None,
                            input_dim=None, output_dim=None,
                            timesteps=None, training=None):
    """Apply `y . w + b` for every temporal slice y of x.

    # Arguments
        x: input tensor.
        w: weight matrix.
        b: optional bias vector.
        dropout: wether to apply dropout (same dropout mask
            for every temporal slice of the input).
        input_dim: integer; optional dimensionality of the input.
        output_dim: integer; optional dimensionality of the output.
        timesteps: integer; optional number of timesteps.
        training: training phase tensor or boolean.

    # Returns
        Output tensor.
    """
    if not input_dim:
        input_dim = K.shape(x)[2]
    if not timesteps:
        timesteps = K.shape(x)[1]
    if not output_dim:
        output_dim = K.int_shape(w)[1]

    if dropout is not None and 0. < dropout < 1.:
        # apply the same dropout pattern at every timestep
        ones = K.ones_like(K.reshape(x[:, 0, :], (-1, input_dim)))
        dropout_matrix = K.dropout(ones, dropout)
        expanded_dropout_matrix = K.repeat(dropout_matrix, timesteps)
        x = K.in_train_phase(x * expanded_dropout_matrix, x, training=training)

    # collapse time dimension and batch dimension together
    x = K.reshape(x, (-1, input_dim))
    x = K.dot(x, w)
    if b is not None:
        x = K.bias_add(x, b)
    # reshape to 3D tensor
    if K.backend() == 'tensorflow':
        x = K.reshape(x, K.stack([-1, timesteps, output_dim]))
        x.set_shape([None, None, output_dim])
    else:
        x = K.reshape(x, (-1, timesteps, output_dim))
    return x 

def call(self, inputs, **kwargs):
        if self.axis == 1:
            # If channels first, force it to be channels last for these ops
            inputs = K.permute_dimensions(inputs, [0, 2, 3, 1])

        q, k, v = tf.split(inputs, [self.depth_k, self.depth_k, self.depth_v], axis=-1)

        q = self.split_heads_2d(q)
        k = self.split_heads_2d(k)
        v = self.split_heads_2d(v)

        # scale query
        depth_k_heads = self.depth_k / self.num_heads
        q *= (depth_k_heads ** -0.5)

        # [Batch, num_heads, height * width, depth_k or depth_v] if axis == -1
        qk_shape = [self._batch, self.num_heads, self._height * self._width, self.depth_k // self.num_heads]
        v_shape = [self._batch, self.num_heads, self._height * self._width, self.depth_v // self.num_heads]
        flat_q = K.reshape(q, K.stack(qk_shape))
        flat_k = K.reshape(k, K.stack(qk_shape))
        flat_v = K.reshape(v, K.stack(v_shape))

        # [Batch, num_heads, HW, HW]
        logits = tf.matmul(flat_q, flat_k, transpose_b=True)

        # Apply relative encodings
        if self.relative:
            h_rel_logits, w_rel_logits = self.relative_logits(q)
            logits += h_rel_logits
            logits += w_rel_logits

        weights = K.softmax(logits, axis=-1)
        attn_out = tf.matmul(weights, flat_v)

        attn_out_shape = [self._batch, self.num_heads, self._height, self._width, self.depth_v // self.num_heads]
        attn_out_shape = K.stack(attn_out_shape)
        attn_out = K.reshape(attn_out, attn_out_shape)
        attn_out = self.combine_heads_2d(attn_out)
        # [batch, height, width, depth_v]

        if self.axis == 1:
            # return to [batch, depth_v, height, width] for channels first
            attn_out = K.permute_dimensions(attn_out, [0, 3, 1, 2])

        return attn_out 

def call(self, inputs):
        X = inputs[0]
        A = inputs[1]

        outputs = []
        output_attn = []
        for head in range(self.attn_heads):
            kernel = self.kernels[head]
            attention_kernel = self.attn_kernels[head]  # Attention kernel a in the paper (2F' x 1)

            # Compute inputs to attention network
            features = K.dot(X, kernel)

            # Compue attention coefficients
            # [[a_1], [a_2]]^T [[Wh_i], [Wh_2]] = [a_1]^T [Wh_i] + [a_2]^T [Wh_j]
            attn_for_self = K.dot(features, attention_kernel[0])    # [a_1]^T [Wh_i]
            attn_for_neighs = K.dot(features, attention_kernel[1])  # [a_2]^T [Wh_j]
            if len(K.int_shape(features)) == 2:
                # Single / mixed mode
                attn_for_neighs_T = K.transpose(attn_for_neighs)
            else:
                # Batch mode
                attn_for_neighs_T = K.permute_dimensions(attn_for_neighs, (0, 2, 1))
            attn_coef = attn_for_self + attn_for_neighs_T
            attn_coef = LeakyReLU(alpha=0.2)(attn_coef)

            # Mask values before activation (Vaswani et al., 2017)
            mask = -10e9 * (1.0 - A)
            attn_coef += mask

            # Apply softmax to get attention coefficients
            attn_coef = K.softmax(attn_coef)
            output_attn.append(attn_coef)

            # Apply dropout to attention coefficients
            attn_coef_drop = Dropout(self.dropout_rate)(attn_coef)

            # Convolution
            features = filter_dot(attn_coef_drop, features)
            if self.use_bias:
                features = K.bias_add(features, self.biases[head])

            # Add output of attention head to final output
            outputs.append(features)

        # Aggregate the heads' output according to the reduction method
        if self.concat_heads:
            output = K.concatenate(outputs)
        else:
            output = K.mean(K.stack(outputs), axis=0)

        output = self.activation(output)

        if self.return_attn_coef:
            return output, output_attn
        else:
            return output 

def call(self, inputs):
        if not isinstance(inputs, list):
            raise ValueError('A merge layer should be called '
                             'on a list of inputs.')
        if self._reshape_required:
            reshaped_inputs = []
            input_ndims = list(map(K.ndim, inputs))
            if None not in input_ndims:
                # If ranks of all inputs are available,
                # we simply expand each of them at axis=1
                # until all of them have the same rank.
                max_ndim = max(input_ndims)
                for x in inputs:
                    x_ndim = K.ndim(x)
                    for _ in range(max_ndim - x_ndim):
                        x = K.expand_dims(x, 1)
                    reshaped_inputs.append(x)
                return self._merge_function(reshaped_inputs)
            else:
                # Transpose all inputs so that batch size is the last dimension.
                # (batch_size, dim1, dim2, ... ) -> (dim1, dim2, ... , batch_size)
                transposed = False
                for x in inputs:
                    x_ndim = K.ndim(x)
                    if x_ndim is None:
                        x_shape = K.shape(x)
                        batch_size = x_shape[0]
                        new_shape = K.concatenate(
                            [x_shape[1:], K.expand_dims(batch_size)])
                        x_transposed = K.reshape(x, K.stack(
                            [batch_size, K.prod(x_shape[1:])]))
                        x_transposed = K.permute_dimensions(x_transposed, (1, 0))
                        x_transposed = K.reshape(x_transposed, new_shape)
                        reshaped_inputs.append(x_transposed)
                        transposed = True
                    elif x_ndim > 1:
                        dims = list(range(1, x_ndim)) + [0]
                        reshaped_inputs.append(K.permute_dimensions(x, dims))
                        transposed = True
                    else:
                        # We don't transpose inputs if they are 1D vectors or scalars.
                        reshaped_inputs.append(x)
                y = self._merge_function(reshaped_inputs)
                y_ndim = K.ndim(y)
                if transposed:
                    # If inputs have been transposed, we have to transpose the output too.
                    if y_ndim is None:
                        y_shape = K.shape(y)
                        y_ndim = K.shape(y_shape)[0]
                        batch_size = y_shape[y_ndim - 1]
                        new_shape = K.concatenate(
                            [K.expand_dims(batch_size), y_shape[:y_ndim - 1]])
                        y = K.reshape(y, (-1, batch_size))
                        y = K.permute_dimensions(y, (1, 0))
                        y = K.reshape(y, new_shape)
                    elif y_ndim > 1:
                        dims = [y_ndim - 1] + list(range(y_ndim - 1))
                        y = K.permute_dimensions(y, dims)
                return y
        else:
            return self._merge_function(inputs) 

def call(self, inputs, **kwargs):
        input_shape = K.int_shape(inputs)
        tensor_input_shape = K.shape(inputs)

        # Prepare broadcasting shape.
        reduction_axes = list(range(len(input_shape)))
        del reduction_axes[self.axis]
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
        broadcast_shape.insert(1, self.groups)

        reshape_group_shape = K.shape(inputs)
        group_axes = [reshape_group_shape[i] for i in range(len(input_shape))]
        group_axes[self.axis] = input_shape[self.axis] // self.groups
        group_axes.insert(1, self.groups)

        # reshape inputs to new group shape
        group_shape = [group_axes[0], self.groups] + group_axes[2:]
        group_shape = K.stack(group_shape)
        inputs = K.reshape(inputs, group_shape)

        group_reduction_axes = list(range(len(group_axes)))
        group_reduction_axes = group_reduction_axes[2:]

        mean = K.mean(inputs, axis=group_reduction_axes, keepdims=True)
        variance = K.var(inputs, axis=group_reduction_axes, keepdims=True)

        inputs = (inputs - mean) / (K.sqrt(variance + self.epsilon))

        # prepare broadcast shape
        inputs = K.reshape(inputs, group_shape)
        outputs = inputs

        # In this case we must explicitly broadcast all parameters.
        if self.scale:
            broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
            outputs = outputs * broadcast_gamma

        if self.center:
            broadcast_beta = K.reshape(self.beta, broadcast_shape)
            outputs = outputs + broadcast_beta

        outputs = K.reshape(outputs, tensor_input_shape)

        return outputs 

def boxes_from_deltas(pred_box_delta, config):
    """
    Converts prediction deltas to bounding boxes
    
    Arguments:
        pred_box_delta {[type]} -- tensor of deltas
        config {[type]} -- hyperparameter dict
    
    Returns:
        [type] -- tensor of bounding boxes
    """



    # Keras backend allows no unstacking

    delta_x = pred_box_delta[:, :, 0]
    delta_y = pred_box_delta[:, :, 1]
    delta_w = pred_box_delta[:, :, 2]
    delta_h = pred_box_delta[:, :, 3]

    # get the coordinates and sizes of the anchor boxes from config

    anchor_x = config.ANCHOR_BOX[:, 0]
    anchor_y = config.ANCHOR_BOX[:, 1]
    anchor_w = config.ANCHOR_BOX[:, 2]
    anchor_h = config.ANCHOR_BOX[:, 3]

    # as we only predict the deltas, we need to transform the anchor box values before computing the loss

    box_center_x = K.identity(
        anchor_x + delta_x * anchor_w)
    box_center_y = K.identity(
        anchor_y + delta_y * anchor_h)
    box_width = K.identity(
        anchor_w * safe_exp(delta_w, config.EXP_THRESH))
    box_height = K.identity(
        anchor_h * safe_exp(delta_h, config.EXP_THRESH))

    # tranform into a real box with four coordinates

    xmins, ymins, xmaxs, ymaxs = bbox_transform([box_center_x, box_center_y, box_width, box_height])

    # trim boxes if predicted outside

    xmins = K.minimum(
        K.maximum(0.0, xmins), config.IMAGE_WIDTH - 1.0)
    ymins = K.minimum(
        K.maximum(0.0, ymins), config.IMAGE_HEIGHT - 1.0)
    xmaxs = K.maximum(
        K.minimum(config.IMAGE_WIDTH - 1.0, xmaxs), 0.0)
    ymaxs = K.maximum(
        K.minimum(config.IMAGE_HEIGHT - 1.0, ymaxs), 0.0)

    det_boxes = K.permute_dimensions(
        K.stack(bbox_transform_inv([xmins, ymins, xmaxs, ymaxs])),
        (1, 2, 0)
    )
    
    return (det_boxes) 

def boxes_from_deltas_np(pred_box_delta, config):

    """
    Converts prediction deltas to bounding boxes, but in numpy
    
    Arguments:
        pred_box_delta {[type]} -- tensor of deltas
        config {[type]} -- hyperparameter dict
    
    Returns:
        [type] -- tensor of bounding boxes
    """


    # Keras backend allows no unstacking

    delta_x = pred_box_delta[:, :, 0]
    delta_y = pred_box_delta[:, :, 1]
    delta_w = pred_box_delta[:, :, 2]
    delta_h = pred_box_delta[:, :, 3]

    # get the coordinates and sizes of the anchor boxes from config

    anchor_x = config.ANCHOR_BOX[:, 0]
    anchor_y = config.ANCHOR_BOX[:, 1]
    anchor_w = config.ANCHOR_BOX[:, 2]
    anchor_h = config.ANCHOR_BOX[:, 3]

    # as we only predict the deltas, we need to transform the anchor box values before computing the loss

    box_center_x = anchor_x + delta_x * anchor_w
    box_center_y = anchor_y + delta_y * anchor_h
    box_width = anchor_w * safe_exp_np(delta_w, config.EXP_THRESH)
    box_height = anchor_h * safe_exp_np(delta_h, config.EXP_THRESH)

    # tranform into a real box with four coordinates

    xmins, ymins, xmaxs, ymaxs = bbox_transform([box_center_x, box_center_y, box_width, box_height])

    # trim boxes if predicted outside

    xmins = np.minimum(
        np.maximum(0.0, xmins), config.IMAGE_WIDTH - 1.0)
    ymins = np.minimum(
        np.maximum(0.0, ymins), config.IMAGE_HEIGHT - 1.0)
    xmaxs = np.maximum(
        np.minimum(config.IMAGE_WIDTH - 1.0, xmaxs), 0.0)
    ymaxs = np.maximum(
        np.minimum(config.IMAGE_HEIGHT - 1.0, ymaxs), 0.0)

    det_boxes = np.transpose(
        np.stack(bbox_transform_inv([xmins, ymins, xmaxs, ymaxs])),
        (1, 2, 0)
    )

    return (det_boxes) 

def _compute_social_tensor(self, grid_t, prev_h_t, config):
        """Compute $H_t_i(m, n, :)$.

        this function implementation is same as  getSocialTensor() function.

        :param grid_t: (batch_size, max_n_peds, max_n_peds, grid_side ** 2)
            which is (batch_index, self_pid, other_pid, grid_index).
        :param prev_h_t: (batch_size, max_n_peds, lstm_state_dim)
        :return: H_t (batch_size, max_n_peds, (grid_side ** 2) * lstm_state_dim)
        """
        H_t = []

        for i in range(config.max_n_peds):
            # (batch_size, max_n_peds, max_n_peds, grid_side ** 2)
            # => (batch_size, max_n_peds, grid_side ** 2)
            grid_it = Lambda(lambda grid_t: grid_t[:, i, ...])(grid_t)

            # (batch_size, max_n_peds, grid_side **2)
            # => (batch_size, grid_side ** 2, max_n_peds)
            grid_it_T = Permute((2, 1))(grid_it)

            # (batch_size, grid_side ** 2, lstm_state_dim)
            H_it = Lambda(lambda x: K.batch_dot(x[0], x[1]))(
                [grid_it_T, prev_h_t])

            # store to H_t
            H_t.append(H_it)

        # list of (batch_size, grid_side_squared, lstm_state_dim)
        # => (max_n_peds, batch_size, grid_side_squared, lstm_state_dim)
        H_t = Lambda(lambda H_t: K.stack(H_t, axis=0))(H_t)

        # (max_n_peds, batch_size, grid_side_squared, lstm_state_dim)
        # => (batch_size, max_n_peds, grid_side_squared, lstm_state_dim)
        H_t = Lambda(lambda H_t: K.permute_dimensions(H_t, (1, 0, 2, 3)))(H_t)

        # (batch_size, max_n_peds, grid_side_squared, lstm_state_dim)
        # => (batch_size, max_n_peds, grid_side_squared * lstm_state_dim)
        H_t = Reshape(
            (config.max_n_peds,
             config.grid_side_squared * config.lstm_state_dim))(
            H_t)

        return H_t 

def _time_distributed_dense(x, w, b=None, dropout=None,
                            input_dim=None, output_dim=None,
                            timesteps=None, training=None):
    """Apply `y . w + b` for every temporal slice y of x.
    # Arguments
        x: input tensor.
        w: weight matrix.
        b: optional bias vector.
        dropout: wether to apply dropout (same dropout mask
            for every temporal slice of the input).
        input_dim: integer; optional dimensionality of the input.
        output_dim: integer; optional dimensionality of the output.
        timesteps: integer; optional number of timesteps.
        training: training phase tensor or boolean.
    # Returns
        Output tensor.
    """
    if not input_dim:
        input_dim = K.shape(x)[2]
    if not timesteps:
        timesteps = K.shape(x)[1]
    if not output_dim:
        output_dim = K.int_shape(w)[1]

    if dropout is not None and 0. < dropout < 1.:
        # apply the same dropout pattern at every timestep
        ones = K.ones_like(K.reshape(x[:, 0, :], (-1, input_dim)))
        dropout_matrix = K.dropout(ones, dropout)
        expanded_dropout_matrix = K.repeat(dropout_matrix, timesteps)
        x = K.in_train_phase(x * expanded_dropout_matrix, x, training=training)

    # collapse time dimension and batch dimension together
    x = K.reshape(x, (-1, input_dim))
    x = K.dot(x, w)
    if b is not None:
        x = K.bias_add(x, b)
    # reshape to 3D tensor
    if K.backend() == 'tensorflow':
        x = K.reshape(x, K.stack([-1, timesteps, output_dim]))
        x.set_shape([None, None, output_dim])
    else:
        x = K.reshape(x, (-1, timesteps, output_dim))
    return x 

def call(self, inputs, **kwargs):
        input_shape = K.int_shape(inputs)
        tensor_input_shape = K.shape(inputs)

        # Prepare broadcasting shape.
        reduction_axes = list(range(len(input_shape)))
        del reduction_axes[self.axis]
        broadcast_shape = [1] * len(input_shape)
        broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
        broadcast_shape.insert(1, self.groups)

        reshape_group_shape = K.shape(inputs)
        group_axes = [reshape_group_shape[i] for i in range(len(input_shape))]
        group_axes[self.axis] = input_shape[self.axis] // self.groups
        group_axes.insert(1, self.groups)

        # reshape inputs to new group shape
        group_shape = [group_axes[0], self.groups] + group_axes[2:]
        group_shape = K.stack(group_shape)
        inputs = K.reshape(inputs, group_shape)

        group_reduction_axes = list(range(len(group_axes)))
        mean, variance = KC.moments(inputs, group_reduction_axes[2:],
                                    keep_dims=True)
        inputs = (inputs - mean) / (K.sqrt(variance + self.epsilon))

        # prepare broadcast shape
        inputs = K.reshape(inputs, group_shape)

        outputs = inputs

        # In this case we must explicitly broadcast all parameters.
        if self.scale:
            broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
            outputs = outputs * broadcast_gamma

        if self.center:
            broadcast_beta = K.reshape(self.beta, broadcast_shape)
            outputs = outputs + broadcast_beta

        # finally we reshape the output back to the input shape
        outputs = K.reshape(outputs, tensor_input_shape)

        return outputs 

def step(self, layer_input, states):
        # As a step function MUST return its regular output as the first element in the list of states,
        # we have _ here.
        _, M, weights_read_tm1, weights_write_tm1 = states[:4]

        # reshaping (TODO: figure out how save n-dimensional state) 
        weights_read_tm1 = K.reshape(weights_read_tm1,
                (self.batch_size, self.read_heads, self.n_slots))
        weights_write_tm1 = K.reshape(weights_write_tm1, 
                (self.batch_size, self.write_heads, self.n_slots))

        # We have the old memory M, and a read weighting w_read_tm1 calculated in the last
        # step. This is enough to calculate the read_vector we feed into the controller:
        memory_read_input = K.concatenate([self._read_from_memory(weights_read_tm1[:,i], M) for 
                i in range(self.read_heads)])

        # Now feed the controller and let it run a single step, implemented by calling the step function directly,
        # which we have to provide with the actual input from outside, the information we've read an the states which
        # are relevant to the controller.
        controller_output = self._run_controller(layer_input, memory_read_input)

        # We take the big chunk of unactivated controller output and subdivide it into actual output, reading and
        # writing instructions. Also specific activions for each parameter are applied.
        ntm_output, controller_instructions_read, controller_instructions_write = \
                self._split_and_apply_activations(controller_output)


        # Now we want to write to the memory for each head. We have to be carefull about concurrency, otherwise there is
        # a chance the write heads will interact with each other in unintended ways!
        # We first calculate all the weights, then perform all the erasing and only after that the adding is done.
        # addressing:
        weights_write = []
        for i in range(self.write_heads):
            write_head = controller_instructions_write[i]
            old_weight_vector = weights_write_tm1[:,i]
            weight_vector = self._get_weight_vector(M, old_weight_vector, *tuple(write_head[:5]))
            weights_write.append(weight_vector)
        # erasing:
        for i in range(self.write_heads):
            M = self._write_to_memory_erase(M, weights_write[i], controller_instructions_write[i][5])
        # adding:
        for i in range(self.write_heads):
            M = self._write_to_memory_add(M, weights_write[i], controller_instructions_write[i][6])

        # Only one thing left until this step is complete: Calculate the read weights we save in the state and use next
        # round:
        # As reading is side-effect-free, we dont have to worry about concurrency.
        weights_read = []
        for i in range(self.read_heads):
            read_head = controller_instructions_read[i]
            old_weight_vector = weights_read_tm1[:,i]
            weight_vector = self._get_weight_vector(M, old_weight_vector, *read_head)
            weights_read.append(weight_vector)

        # M = tf.Print(M, [K.mean(M), K.max(M), K.min(M)], message="Memory overview")
        # Now lets pack up the state in a list and call it a day.
        return ntm_output, [ntm_output, M, K.stack(weights_read, axis=1), K.stack(weights_write, axis=1)] 

def call(self, inputs):
        if self.data_format == 'channels_first':
            # Pad the vertical
            if self.padding[0][0] > 0:
                top_slice = K.stack([inputs[:, :, 0]] * self.padding[0][0], axis=2)
            else:
                top_slice = inputs[:, :, slice(0, 0)]
            if self.padding[0][1] > 0:
                bottom_slice = K.stack([inputs[:, :, -1]] * self.padding[0][1], axis=2)
            else:
                bottom_slice = inputs[:, :, slice(0, 0)]
            outputs = K.concatenate([top_slice, inputs, bottom_slice], axis=2)
            # Pad the horizontal
            if self.padding[1][0] > 0:
                left_slice = K.stack([outputs[:, :, :, 0]] * self.padding[1][0], axis=3)
            else:
                left_slice = outputs[:, :, :, slice(0, 0)]
            if self.padding[1][1] > 0:
                right_slice = K.stack([outputs[:, :, :, -1]] * self.padding[1][1], axis=3)
            else:
                right_slice = outputs[:, :, :, slice(0, 0)]
            outputs = K.concatenate([left_slice, outputs, right_slice], axis=3)
        else:
            # Pad the vertical
            if self.padding[0][0] > 0:
                top_slice = K.stack([inputs[:, 0]] * self.padding[0][0], axis=1)
            else:
                top_slice = inputs[:, slice(0, 0)]
            if self.padding[0][1] > 0:
                bottom_slice = K.stack([inputs[:, -1]] * self.padding[0][1], axis=1)
            else:
                bottom_slice = inputs[:, slice(0, 0)]
            outputs = K.concatenate([top_slice, inputs, bottom_slice], axis=1)
            # Pad the horizontal
            if self.padding[1][0] > 0:
                left_slice = K.stack([outputs[:, :, 0]] * self.padding[1][0], axis=2)
            else:
                left_slice = outputs[:, :, slice(0, 0)]
            if self.padding[1][1] > 0:
                right_slice = K.stack([outputs[:, :, -1]] * self.padding[1][1], axis=2)
            else:
                right_slice = outputs[:, :, slice(0, 0)]
            outputs = K.concatenate([left_slice, outputs, right_slice], axis=2)
        return outputs 

def call(self, X):
        ### here X is the entire mat with [batch, T, D]
        n_frames = X.get_shape().as_list()[1]
        z_list = []

        ## before iteration, we do the first bilinear fusion
        lambda_0 = K.reshape(K.abs(self.core_diag_list[0]), [1, 1, -1])
        z = lambda_0 * X * X
        z = self.act_fun_out(z)
        z_list.append(z)

        for ii in range(self.n_recur):

            ## (1) obtain x_{t+1} = U_tx_t + b_t 
            X = K.dot(X, self.factorInMats[ii]) + self.factorInBias[ii]
            X = self.act_fun_in(X)
            lambda_i = K.reshape(K.abs(self.core_diag_list[ii+1]), [1, 1, -1])
            z = lambda_i * X * X
            zout = K.dot(z, self.factorOutMats[ii]) + self.factorOutBias[ii]
            zout = self.act_fun_out(zout)
            z_list.append(zout)

        if self.out_fusion_type == 'mean':
            ## compute the mean value 
            z_list_tensor = K.stack(z_list, axis=-2)
            out = K.mean(z_list_tensor, axis=-2)
        
        elif self.out_fusion_type == 'concate':
            ## concatente features
            out = K.concatenate(z_list, axis=-1)

        else:
            print('[ERROR] no such fusion method. Program terminates')
            sys.exit()


        self.out_dim = out.get_shape().as_list()[-1]
        ### now out is [batch, T, out_dim], we do temporal local pooling
        #### zero padding
        # out = ZeroPadding1D((self.time_window_size//2))(out)
        # W = tf.reshape(self.conv_filter, [1, -1, 1]) # [1, |Nt|, 1]
        # out_pool_list = [ K.sum(out[:, i:i+self.time_window_size, :]*W, axis=1)
        #                  for i in range(0,n_frames, self.stride) ]

        # out_pool = K.stack(out_pool_list,axis=1)
        out_pool = AveragePooling1D(self.stride)(out)

        return out_pool 

def call(self, X):
        ### here X is the entire mat with [batch, T, D]
        n_frames = X.get_shape().as_list()[1]
        z_list = []

        ## before iteration, we do the first bilinear fusion
        lambda_0 = K.reshape(K.abs(self.core_diag_list[0]), [1, 1, -1])
        z = lambda_0 * X * X
        z = self.act_fun_out(z)
        z_list.append(z)

        for ii in range(self.n_recur):
            ## (1) obtain x_{t+1} = U_tx_t + b_t 
            # U = QRDecompose()(self.factorInMats[ii])
            U = self.factorInMats[ii]
            if self.use_bias:
                X1 = K.dot(X, U) + self.factorInBias[ii]
            else:
                X1 = K.dot(X, U)

            X1 = self.act_fun_in(X1)
            lambda_i = K.reshape(K.abs(self.core_diag_list[ii+1]), [1, 1, -1])
            z = lambda_i * X1 * X1

            # W = QRDecompose()(self.factorOutMats[ii])
            W = self.factorOutMats[ii]

            if self.use_bias:
                zout = K.dot(z, W) + self.factorOutBias[ii]
            else:
                zout = K.dot(z, W)

            zout = self.act_fun_out(zout)
            z_list.append(zout)

        if self.out_fusion_type == 'sum':
            # ## compute the mean value 
            z_list_tensor = K.stack(z_list, axis=-1)
            out = K.sum(z_list_tensor, axis=-1)

        elif self.out_fusion_type == 'concate':
            ## concatente features
            out = K.concatenate(z_list, axis=-1)
            
        self.out_dim = out.get_shape().as_list()[-1]

        ## now out is [batch, T, out_dim], we do temporal local pooling
        ### zero padding
        # out = ZeroPadding1D((self.time_window_size//2))(out)
        # W = tf.reshape(self.conv_filter, [1, -1, 1]) # [1, |Nt|, 1]
        # out_pool_list = [ K.sum(out[:, i:i+self.time_window_size, :]*W, axis=1)
        #                  for i in range(0,n_frames, self.stride)  ]

        # out_pool = K.stack(out_pool_list,axis=1)
        out_pool = AveragePooling1D(self.stride)(out) 
        # out_pool = MaxPooling1D(self.stride)(out) 

        return out_pool 

def call(self, X):
        ### here X is the entire mat with [batch, T, D]
        n_frames = X.get_shape().as_list()[1]
        # self.E = K.l2_normalize(self.E, axis=0)
        # self.F = K.l2_normalize(self.F, axis=0)
        # self.E_softmax = K.softmax(self.E, axis=0)
        # self.F_softmax = K.softmax(self.F, axis=0)
        
        # apply uni-norm, non-negative constraint
        #self.E *= K.cast(K.greater_equal(self.E, 0), float) # non negative constraint
        #
        #self.E = self.E / (K.epsilon() + K.sqrt(K.sum(K.square(self.E),
        #                                       axis=0,
        #                                       keepdims=True)))
        #self.F *= K.cast(K.greater_equal(self.F, 0.), float) # non negative constraint
        #
        #self.F = self.F / (K.epsilon() + K.sqrt(K.sum(K.square(self.F),
        #                                       axis=0,
        #                                       keepdims=True)))


        z1 = K.dot(X, self.E  )+K.reshape(self.bx, [1,1,-1])
        z2 = K.dot(X, self.F  )+K.reshape(self.by, [1,1,-1])

        if self.activation == 'tanh':
            z1 = K.tanh(z1)
            z2 = K.tanh(z2)


        # outer product
        z1 = K.expand_dims(z1, axis=-1)
        z2 = K.expand_dims(z2, axis=-2)
        z = tf.matmul(z1, z2)
        z = self.G*K.reshape(z, [-1, n_frames, self.n_basis**2])
        
       

        ## pooling to reduce time sequence
        if self.out_fusion_type == 'avg':
            out_pool = AveragePooling1D(pool_size=self.stride,
                                        strides=self.stride,
                                        padding='same')(z) 
        elif self.out_fusion_type == 'max':
            out_pool = MaxPooling1D(self.stride)(z) 
        elif self.out_fusion_type == 'w-sum':
            ### now out is [batch, T, out_dim], we do temporal local pooling
            #### zero padding
            out = ZeroPadding1D((self.time_window_size//2))(z)

            W = tf.reshape(self.conv_filter, [1, -1, 1]) # [1, |Nt|, 1]
            out_pool_list = [ K.sum(out[:, i:i+self.time_window_size, :]*W, axis=1)
                             for i in range(0,n_frames, self.stride)  ]

            out_pool = K.stack(out_pool_list,axis=1)
        elif self.out_fusion_type == 'linearproj':
            out_pool = Conv1D(self.out_dim, 1, strides=self.stride, padding='same')(z)


        return out_pool 

def call(self, X):
        ### here X is the entire mat with [batch, T, D]
        n_frames = X.get_shape().as_list()[1]
        # self.E = K.l2_normalize(self.E, axis=0)
        # self.F = K.l2_normalize(self.F, axis=0)
        # self.E_softmax = K.softmax(self.E, axis=0)
        # self.F_softmax = K.softmax(self.F, axis=0)

        z1 = self.E_layer(X)
        z2 = self.F_layer(X)

        # if self.activation == 'tanh':
        #     z1 = K.tanh(z1)
        #     z2 = K.tanh(z2)


        # outer product
        z1 = K.expand_dims(z1, axis=-1)
        z2 = K.expand_dims(z2, axis=-2)
        z = tf.matmul(z1, z2)
        z = K.reshape(z, [-1, n_frames, self.n_basis**2])
        

        ## use power and l2 normalization
        if self.use_normalization:
            z = SqrtAcfun(theta=1e-3)(z)
            z = K.l2_normalize(z, axis=-1)


        ## pooling to reduce time sequence
        if self.out_fusion_type == 'avg':
            out_pool = AveragePooling1D(pool_size=self.stride,
                                        strides=self.stride,
                                        padding='same')(z) 
        elif self.out_fusion_type == 'max':
            out_pool = MaxPooling1D(self.stride)(z) 
        elif self.out_fusion_type == 'w-sum':
            ### now out is [batch, T, out_dim], we do temporal local pooling
            #### zero padding
            out = ZeroPadding1D((self.time_window_size//2))(z)

            W = tf.reshape(self.conv_filter, [1, -1, 1]) # [1, |Nt|, 1]
            out_pool_list = [ K.sum(out[:, i:i+self.time_window_size, :]*W, axis=1)
                             for i in range(0,n_frames, self.stride)  ]

            out_pool = K.stack(out_pool_list,axis=1)
        elif self.out_fusion_type == 'linearproj':
            out_pool = Conv1D(self.out_dim, 1, strides=self.stride, padding='same')(z)


        return out_pool 

def _time_distributed_dense(x, w, b=None, dropout=None,
                            input_dim=None, output_dim=None,
                            timesteps=None, training=None):
    """Apply `y . w + b` for every temporal slice y of x.
    # Arguments
        x: input tensor.
        w: weight matrix.
        b: optional bias vector.
        dropout: wether to apply dropout (same dropout mask
            for every temporal slice of the input).
        input_dim: integer; optional dimensionality of the input.
        output_dim: integer; optional dimensionality of the output.
        timesteps: integer; optional number of timesteps.
        training: training phase tensor or boolean.
    # Returns
        Output tensor.
    """
    if not input_dim:
        input_dim = K.shape(x)[2]
    if not timesteps:
        timesteps = K.shape(x)[1]
    if not output_dim:
        output_dim = K.shape(w)[1]

    if dropout is not None and 0. < dropout < 1.:
        # apply the same dropout pattern at every timestep
        ones = K.ones_like(K.reshape(x[:, 0, :], (-1, input_dim)))
        dropout_matrix = K.dropout(ones, dropout)
        expanded_dropout_matrix = K.repeat(dropout_matrix, timesteps)
        x = K.in_train_phase(x * expanded_dropout_matrix, x, training=training)

    # collapse time dimension and batch dimension together
    x = K.reshape(x, (-1, input_dim))
    x = K.dot(x, w)
    if b is not None:
        x = K.bias_add(x, b)
    # reshape to 3D tensor
    if K.backend() == 'tensorflow':
        x = K.reshape(x, K.stack([-1, timesteps, output_dim]))
        x.set_shape([None, None, output_dim])
    else:
        x = K.reshape(x, (-1, timesteps, output_dim))
    return x 

def layered_weighted_bce(y_true, y_pred, weights):
    """Binary cross-entropy function with different weights for mask channels.

    Arguments:
    ----------
    y_true (tensor): passed silently by Keras during model training.
    y_pred (tensor): passed silently by Keras during model training.
    weights (list-like): Weights to assign to mask foreground pixels for each
        channel in the 3rd axis of the mask.

    Returns:
    --------
    The binary cross-entropy loss function output multiplied by a weighting
    mask.

    Usage:
    ------
    See implementation instructions for `weighted_bce`.

    This loss function is intended to allow different weighting of different
    segmentation outputs - for example, if a model outputs a 3D image mask,
    where the first channel corresponds to foreground objects and the second
    channel corresponds to object edges. `weights` must be a list of length
    equal to the depth of the output mask. The output mask's "z-axis"
    corresponding to the mask channel must be the third axis in the output
    array.

    """
    weight_mask = K.ones_like(y_true)
    submask_list = []
    for i in range(len(weights)):
        class_two = K.equal(y_true[:, :, :, i], weight_mask[:, :, :, i])
        class_two = K.cast(class_two, 'float32')
        if weights[i] < 1:
            class_two = class_two*(1-weights[i])
            layer_mask = weight_mask[:, :, :, i] - class_two
        elif weights[i] > 1:
            class_two = class_two*(weights[i]-1)
            layer_mask = weight_mask[:, :, :, i] + class_two
        else:
            layer_mask = weight_mask[:, :, :, i]
        submask_list.append(layer_mask)
    final_mask = K.stack(submask_list, axis=-1)
    return K.binary_crossentropy(y_pred, y_true) * final_mask 

def _time_distributed_dense(x, w, b=None, dropout=None,
                            input_dim=None, output_dim=None,
                            timesteps=None, training=None):
    """Apply `y . w + b` for every temporal slice y of x.
    # Arguments
        x: input tensor.
        w: weight matrix.
        b: optional bias vector.
        dropout: wether to apply dropout (same dropout mask
            for every temporal slice of the input).
        input_dim: integer; optional dimensionality of the input.
        output_dim: integer; optional dimensionality of the output.
        timesteps: integer; optional number of timesteps.
        training: training phase tensor or boolean.
    # Returns
        Output tensor.
    """
    if not input_dim:
        input_dim = K.shape(x)[2]
    if not timesteps:
        timesteps = K.shape(x)[1]
    if not output_dim:
        output_dim = K.shape(w)[1]

    if dropout is not None and 0. < dropout < 1.:
        # apply the same dropout pattern at every timestep
        ones = K.ones_like(K.reshape(x[:, 0, :], (-1, input_dim)))
        dropout_matrix = K.dropout(ones, dropout)
        expanded_dropout_matrix = K.repeat(dropout_matrix, timesteps)
        x = K.in_train_phase(x * expanded_dropout_matrix, x, training=training)

    # collapse time dimension and batch dimension together
    x = K.reshape(x, (-1, input_dim))
    x = K.dot(x, w)
    if b is not None:
        x = K.bias_add(x, b)
    # reshape to 3D tensor
    if K.backend() == 'tensorflow':
        x = K.reshape(x, K.stack([-1, timesteps, output_dim]))
        x.set_shape([None, None, output_dim])
    else:
        x = K.reshape(x, (-1, timesteps, output_dim))
    return x