How to use the deeppavlov.models.coreference_resolution.new_model.custom_layers.shape function in deeppavlov

def get_slow_antecedent_scores(self, top_span_emb, top_antecedents, top_antecedent_emb, top_antecedent_offsets,
                                   top_span_speaker_ids, genre_emb):
        k = custom_layers.shape(top_span_emb, 0)
        c = custom_layers.shape(top_antecedents, 1)

        feature_emb_list = []

        if self.use_metadata:
            top_antecedent_speaker_ids = tf.gather(top_span_speaker_ids, top_antecedents)  # [k, c]
            same_speaker = tf.equal(tf.expand_dims(top_span_speaker_ids, 1), top_antecedent_speaker_ids)  # [k, c]
            speaker_pair_emb = tf.gather(tf.get_variable("same_speaker_emb", [2, self.feature_size]),
                                         tf.to_int32(same_speaker))  # [k, c, emb]
            feature_emb_list.append(speaker_pair_emb)

            tiled_genre_emb = tf.tile(tf.expand_dims(tf.expand_dims(genre_emb, 0), 0), [k, c, 1])  # [k, c, emb]
            feature_emb_list.append(tiled_genre_emb)

        if self.use_features:
            antecedent_distance_buckets = self.bucket_distance(top_antecedent_offsets)  # [k, c]
            antecedent_distance_emb = tf.gather(

context_emb_list.append(aggregated_char_emb)
            head_emb_list.append(aggregated_char_emb)

        # if not self.lm_file:
        #     elmo_module = hub.Module("https://tfhub.dev/google/elmo/2")
        #     lm_embeddings = elmo_module(
        #         inputs={"tokens": tokens, "sequence_len": text_len},
        #         signature="tokens", as_dict=True)
        #     word_emb = lm_embeddings["word_emb"]  # [num_sentences, max_sentence_length, 512]
        #     lm_emb = tf.stack([tf.concat([word_emb, word_emb], -1),
        #                        lm_embeddings["lstm_outputs1"],
        #                        lm_embeddings["lstm_outputs2"]], -1)  # [num_sentences, max_sentence_length, 1024, 3]

        lm_emb_size = custom_layers.shape(lm_emb, 2)
        lm_num_layers = custom_layers.shape(lm_emb, 3)
        with tf.variable_scope("lm_aggregation"):
            self.lm_weights = tf.nn.softmax(
                tf.get_variable("lm_scores", [lm_num_layers], initializer=tf.constant_initializer(0.0)))
            self.lm_scaling = tf.get_variable("lm_scaling", [], initializer=tf.constant_initializer(1.0))
        flattened_lm_emb = tf.reshape(lm_emb, [num_sentences * max_sentence_length * lm_emb_size, lm_num_layers])
        flattened_aggregated_lm_emb = tf.matmul(flattened_lm_emb, tf.expand_dims(self.lm_weights, 1))
        # [num_sentences * max_sentence_length * emb, 1]

        aggregated_lm_emb = tf.reshape(flattened_aggregated_lm_emb, [num_sentences, max_sentence_length, lm_emb_size])
        aggregated_lm_emb *= self.lm_scaling
        context_emb_list.append(aggregated_lm_emb)

        context_emb = tf.concat(context_emb_list, 2)  # [num_sentences, max_sentence_length, emb]
        head_emb = tf.concat(head_emb_list, 2)  # [num_sentences, max_sentence_length, emb]
        context_emb = tf.nn.dropout(context_emb, self.lexical_dropout)  # [num_sentences, max_sentence_length, emb]
        head_emb = tf.nn.dropout(head_emb, self.lexical_dropout)  # [num_sentences, max_sentence_length, emb]

self.lexical_dropout = self.get_dropout(self.lexical_dropout_rate, is_training)
        self.lstm_dropout = self.get_dropout(self.lstm_dropout_rate, is_training)

        num_sentences = tf.shape(context_word_emb)[0]
        max_sentence_length = tf.shape(context_word_emb)[1]

        context_emb_list = [context_word_emb]
        head_emb_list = [head_word_emb]

        if self.char_embedding_size > 0:
            char_emb = tf.gather(
                tf.get_variable("char_embeddings", [len(self.char_dict), self.char_embedding_size]),
                char_index)  # [num_sentences, max_sentence_length, max_word_length, emb]
            flattened_char_emb = tf.reshape(char_emb,
                                            [num_sentences * max_sentence_length,
                                             custom_layers.shape(char_emb, 2),
                                             custom_layers.shape(char_emb, 3)])
            # [num_sentences * max_sentence_length, max_word_length, emb]

            flattened_aggregated_char_emb = custom_layers.cnn(flattened_char_emb, self.filter_widths, self.filter_size)
            # [num_sentences * max_sentence_length, emb]

            aggregated_char_emb = tf.reshape(flattened_aggregated_char_emb, [num_sentences, max_sentence_length,
                                                                             custom_layers.shape(
                                                                                 flattened_aggregated_char_emb, 1)])
            # [num_sentences, max_sentence_length, emb]

            context_emb_list.append(aggregated_char_emb)
            head_emb_list.append(aggregated_char_emb)

        # if not self.lm_file:
        #     elmo_module = hub.Module("https://tfhub.dev/google/elmo/2")

def flatten_emb_by_sentence(emb, text_len_mask):
        num_sentences = tf.shape(emb)[0]
        max_sentence_length = tf.shape(emb)[1]

        emb_rank = len(emb.get_shape())
        if emb_rank == 2:
            flattened_emb = tf.reshape(emb, [num_sentences * max_sentence_length])
        elif emb_rank == 3:
            flattened_emb = tf.reshape(emb, [num_sentences * max_sentence_length, custom_layers.shape(emb, 2)])
        else:
            raise ValueError("Unsupported rank: {}".format(emb_rank))
        return tf.boolean_mask(flattened_emb, tf.reshape(text_len_mask, [num_sentences * max_sentence_length]))

flattened_aggregated_lm_emb = tf.matmul(flattened_lm_emb, tf.expand_dims(self.lm_weights, 1))
        # [num_sentences * max_sentence_length * emb, 1]

        aggregated_lm_emb = tf.reshape(flattened_aggregated_lm_emb, [num_sentences, max_sentence_length, lm_emb_size])
        aggregated_lm_emb *= self.lm_scaling
        context_emb_list.append(aggregated_lm_emb)

        context_emb = tf.concat(context_emb_list, 2)  # [num_sentences, max_sentence_length, emb]
        head_emb = tf.concat(head_emb_list, 2)  # [num_sentences, max_sentence_length, emb]
        context_emb = tf.nn.dropout(context_emb, self.lexical_dropout)  # [num_sentences, max_sentence_length, emb]
        head_emb = tf.nn.dropout(head_emb, self.lexical_dropout)  # [num_sentences, max_sentence_length, emb]

        text_len_mask = tf.sequence_mask(text_len, maxlen=max_sentence_length)  # [num_sentence, max_sentence_length]

        context_outputs = self.lstm_contextualize(context_emb, text_len, text_len_mask)  # [num_words, emb]
        num_words = custom_layers.shape(context_outputs, 0)

        genre_emb = tf.gather(tf.get_variable("genre_embeddings", [len(self.genres), self.feature_size]),
                              genre)  # [emb]
        # --------------------------------------------------------------------------------------------------------------
        sentence_indices = tf.tile(tf.expand_dims(tf.range(num_sentences), 1),
                                   [1, max_sentence_length])  # [num_sentences, max_sentence_length]
        flattened_sentence_indices = self.flatten_emb_by_sentence(sentence_indices, text_len_mask)  # [num_words]
        flattened_head_emb = self.flatten_emb_by_sentence(head_emb, text_len_mask)  # [num_words]

        k = tf.to_int32(tf.floor(tf.to_float(tf.shape(context_outputs)[0]) * self.top_span_ratio))
        c = tf.minimum(self.max_top_antecedents, k)

        if self.train_on_gold:
            candidate_cluster_ids = self.get_candidate_labels(gold_starts, gold_ends, gold_starts, gold_ends,
                                                              cluster_ids)  # [num_candidates]
            candidate_span_emb = self.get_span_emb(flattened_head_emb, context_outputs, gold_starts,

top_antecedent_scores = \
                    top_fast_antecedent_scores + self.get_slow_antecedent_scores(top_span_emb,
                                                                                 top_antecedents,
                                                                                 top_antecedent_emb,
                                                                                 top_antecedent_offsets,
                                                                                 top_span_speaker_ids,
                                                                                 genre_emb)  # [k, c]
                top_antecedent_weights = tf.nn.softmax(
                    tf.concat([dummy_scores, top_antecedent_scores], 1))  # [k, c + 1]
                top_antecedent_emb = tf.concat([tf.expand_dims(top_span_emb, 1), top_antecedent_emb], 1)
                # [k, c + 1, emb]
                attended_span_emb = tf.reduce_sum(tf.expand_dims(top_antecedent_weights, 2) * top_antecedent_emb,
                                                  1)  # [k, emb]
                with tf.variable_scope("f"):
                    f = tf.sigmoid(custom_layers.projection(tf.concat([top_span_emb, attended_span_emb], 1),
                                                            custom_layers.shape(top_span_emb, -1)))  # [k, emb]
                    top_span_emb = f * attended_span_emb + (1 - f) * top_span_emb  # [k, emb]

        top_antecedent_scores = tf.concat([dummy_scores, top_antecedent_scores], 1)  # [k, c + 1]

        top_antecedent_cluster_ids = tf.gather(top_span_cluster_ids, top_antecedents)  # [k, c]
        top_antecedent_cluster_ids += tf.to_int32(tf.log(tf.to_float(top_antecedents_mask)))  # [k, c]
        same_cluster_indicator = tf.equal(top_antecedent_cluster_ids, tf.expand_dims(top_span_cluster_ids, 1))  # [k, c]
        non_dummy_indicator = tf.expand_dims(top_span_cluster_ids > 0, 1)  # [k, 1]
        pairwise_labels = tf.logical_and(same_cluster_indicator, non_dummy_indicator)  # [k, c]
        dummy_labels = tf.logical_not(tf.reduce_any(pairwise_labels, 1, keepdims=True))  # [k, 1]
        top_antecedent_labels = tf.concat([dummy_labels, pairwise_labels], 1)  # [k, c + 1]
        loss = self.softmax_loss(top_antecedent_scores, top_antecedent_labels)  # [k]
        loss = tf.reduce_sum(loss)  # []

        return [candidate_mention_scores, top_span_starts, top_span_ends, top_antecedents, top_antecedent_scores], loss

# [num_sentences, max_sentence_length, emb]

            context_emb_list.append(aggregated_char_emb)
            head_emb_list.append(aggregated_char_emb)

        # if not self.lm_file:
        #     elmo_module = hub.Module("https://tfhub.dev/google/elmo/2")
        #     lm_embeddings = elmo_module(
        #         inputs={"tokens": tokens, "sequence_len": text_len},
        #         signature="tokens", as_dict=True)
        #     word_emb = lm_embeddings["word_emb"]  # [num_sentences, max_sentence_length, 512]
        #     lm_emb = tf.stack([tf.concat([word_emb, word_emb], -1),
        #                        lm_embeddings["lstm_outputs1"],
        #                        lm_embeddings["lstm_outputs2"]], -1)  # [num_sentences, max_sentence_length, 1024, 3]

        lm_emb_size = custom_layers.shape(lm_emb, 2)
        lm_num_layers = custom_layers.shape(lm_emb, 3)
        with tf.variable_scope("lm_aggregation"):
            self.lm_weights = tf.nn.softmax(
                tf.get_variable("lm_scores", [lm_num_layers], initializer=tf.constant_initializer(0.0)))
            self.lm_scaling = tf.get_variable("lm_scaling", [], initializer=tf.constant_initializer(1.0))
        flattened_lm_emb = tf.reshape(lm_emb, [num_sentences * max_sentence_length * lm_emb_size, lm_num_layers])
        flattened_aggregated_lm_emb = tf.matmul(flattened_lm_emb, tf.expand_dims(self.lm_weights, 1))
        # [num_sentences * max_sentence_length * emb, 1]

        aggregated_lm_emb = tf.reshape(flattened_aggregated_lm_emb, [num_sentences, max_sentence_length, lm_emb_size])
        aggregated_lm_emb *= self.lm_scaling
        context_emb_list.append(aggregated_lm_emb)

        context_emb = tf.concat(context_emb_list, 2)  # [num_sentences, max_sentence_length, emb]
        head_emb = tf.concat(head_emb_list, 2)  # [num_sentences, max_sentence_length, emb]
        context_emb = tf.nn.dropout(context_emb, self.lexical_dropout)  # [num_sentences, max_sentence_length, emb]

def get_fast_antecedent_scores(self, top_span_emb):
        with tf.variable_scope("src_projection"):
            source_top_span_emb = tf.nn.dropout(
                custom_layers.projection(top_span_emb, custom_layers.shape(top_span_emb, -1)),
                self.dropout)  # [k, emb]
        target_top_span_emb = tf.nn.dropout(top_span_emb, self.dropout)  # [k, emb]
        return tf.matmul(source_top_span_emb, target_top_span_emb, transpose_b=True)  # [k, k]