How to use the mathy.agent.layers.bahdanau_attention.BahdanauAttention function in mathy

self.attention = BahdanauAttention(self.units)
        self.flatten = tf.keras.layers.Flatten(name="flatten")
        self.concat = tf.keras.layers.Concatenate(name=f"embedding_concat")
        self.input_dense = tf.keras.layers.Dense(
            self.units, name="embedding_input", use_bias=False
        )
        self.resnet = ResNetStack(units=units, name="embedding_resnet", num_layers=4)
        self.time_lstm = tf.keras.layers.LSTM(
            units, name="timestep_lstm", return_sequences=True, time_major=True
        )
        self.nodes_lstm = tf.keras.layers.LSTM(
            units, name="nodes_lstm", return_sequences=True, time_major=False
        )
        self.embedding = tf.keras.layers.Dense(units=self.units, name="embedding")
        self.build_feature_columns()
        self.attention = BahdanauAttention(units, name="attention")
        self.dense_h = tf.keras.layers.Dense(units, name=f"lstm_initial_h")
        self.dense_c = tf.keras.layers.Dense(units, name=f"lstm_initial_c")

# to know what the probability of each action is for each node in the expression
        # tree. This is key to allow the model to select which node to apply which action
        # to.
        predict_policy = TimeDistributed(
            MathPolicyDropout(action_size, dropout=dropout_rate)
        )
        policy_predictions = predict_policy(sequence_inputs)
        # Flatten policy logits
        policy_shape = tf.shape(policy_predictions)
        policy_predictions = tf.reshape(
            policy_predictions, [policy_shape[0], -1, 1], name="policy_reshape"
        )

    # Value head
    with tf.compat.v1.variable_scope("love/value_head"):
        attention_context, _ = BahdanauAttention(128)(sequence_inputs, hidden_states)
        value_logits = Dense(1, activation="tanh")(attention_context)

    logits = {"policy": policy_predictions, "value": value_logits}

    # Optimizer (for all tasks)
    optimizer = adam.AdamOptimizer(learning_rate)

    # output histograms for all trainable variables.
    summary_interval = 1000
    global_step = tf.compat.v1.train.get_or_create_global_step()
    with tf.summary.record_if(lambda: tf.math.equal(global_step % summary_interval, 0)):
        for var in tf.compat.v1.trainable_variables():
            tf.compat.v1.summary.histogram(var.name, var)

    with tf.compat.v1.variable_scope("stats"):

# to.
        predict_policy = TimeDistributed(MathPolicyResNet(action_size, dropout=dropout_rate))
        sequence_outputs = predict_policy(sequence_inputs)

        # TODO: I can't get the reshape working to mix the context-features with the sequence
        #  features before doing the softmax prediction above. Without these features I think
        #  the policy output will struggle to improve beyond a certain point. The reasoning is
        #  that the observations are not processed as episode sequences via an RNN of some sort
        #  but are input as distinct observations of the environment state. The sequential features
        #  are accessed by encoding the previous observation into the input features (via the
        #  FEATURE_LAST_FWD_VECTORS and FEATURE_LAST_BWD_VECTORS features). Because of this input
        #  format the context features that describe lifetime and urgency (i.e. current_move,
        #  moves_remaining) cannot be connected to the sequential policy output predictions.
        # TODO: Someone help! ☝️ Thanks! 🙇‍

        attention_context, _ = BahdanauAttention(512)(sequence_outputs, hidden_states)
        network = Concatenate(name="mixed_features")(
            [attention_context, context_inputs]
        )
        value_logits = Dense(1, activation="tanh", name="value_logits")(network)

        # Flatten policy logits
        policy_logits = sequence_outputs
        policy_shape = tf.shape(policy_logits)
        policy_logits = tf.reshape(
            policy_logits, [policy_shape[0], -1, 1], name="policy_reshape"
        )

    logits = {"policy": policy_logits, "value": value_logits}

    # Optimizer (for all tasks)
    optimizer = adam.AdamOptimizer(learning_rate)

units,
            name="nodes_lstm",
            return_sequences=True,
            time_major=False,
            return_state=True,
        )
        self.nodes_lstm_bwd = tf.keras.layers.LSTM(
            units,
            name="nodes_lstm_bwd",
            return_sequences=True,
            time_major=False,
            go_backwards=True,
            return_state=True,
        )
        self.embedding = tf.keras.layers.Dense(units=self.units, name="embedding")
        self.attention = BahdanauAttention(units, name="attention")

policy_head = tf.keras.layers.TimeDistributed(
        MathPolicyDropout(action_size, dropout=dropout_rate, feature_layers=[]),
        name="policy_head",
    )
    policy_predictions = policy_head(lstm_vectors)

    # Value head
    with tf.compat.v1.variable_scope("value_head"):
        attention_context, attention_weights = BahdanauAttention(shared_dense_units)(
            shared_network(lstm_vectors), hidden_states
        )
        value_logits = Dense(1, activation="tanh", name="tanh")(attention_context)

    with tf.compat.v1.variable_scope("auxiliary_heads"):

        aux_attention, aux_attention_weights = BahdanauAttention(shared_dense_units)(
            shared_network(lstm_vectors), hidden_states
        )
        # Node change prediction
        node_ctrl_logits = Dense(1, name="node_ctrl_head")(aux_attention)
        # Grouping error prediction
        grouping_ctrl_logits = Dense(1, name="grouping_ctrl_head")(aux_attention)
        # Group prediction head is an integer value predicting the number
        #  of like-term groups in the observation.
        group_prediction_logits = Dense(1, name="group_prediction_head")(aux_attention)
        # Reward prediction head with 3 class labels (positive, negative, neutral)
        reward_prediction_logits = Dense(3, name="reward_prediction_head")(
            aux_attention
        )

    logits = {
        "policy": policy_predictions,

def __init__(
        self,
        units: int = 128,
        sequence_steps: int = 3,
        burn_in_steps: int = 10,
        **kwargs,
    ):
        super(MathyEmbedding, self).__init__(**kwargs)
        self.units = units
        self.burn_in_steps = burn_in_steps
        self.sequence_steps = sequence_steps
        self.init_rnn_state(self.units)
        self.attention = BahdanauAttention(self.units)
        self.flatten = tf.keras.layers.Flatten(name="flatten")
        self.concat = tf.keras.layers.Concatenate(name=f"embedding_concat")
        self.input_dense = tf.keras.layers.Dense(
            self.units, name="embedding_input", use_bias=False
        )
        self.resnet = ResNetStack(units=units, name="embedding_resnet", num_layers=4)
        self.time_lstm = tf.keras.layers.LSTM(
            units, name="timestep_lstm", return_sequences=True, time_major=True
        )
        self.nodes_lstm = tf.keras.layers.LSTM(
            units,
            name="nodes_lstm",
            return_sequences=True,
            time_major=False,
            return_state=True,
        )

def __init__(self, units: int = 128, **kwargs):
        super(MathEmbedding, self).__init__(**kwargs)
        self.units = units
        self.attention = BahdanauAttention(self.units)
        self.flatten = tf.keras.layers.Flatten(name="flatten")
        self.concat = tf.keras.layers.Concatenate(name=f"embedding_concat")
        self.input_dense = tf.keras.layers.Dense(
            self.units, name="embedding_input", use_bias=False
        )
        self.resnet = ResNetStack(units=units, name="embedding_resnet", num_layers=4)
        self.time_lstm = tf.keras.layers.LSTM(
            units, name="timestep_lstm", return_sequences=True, time_major=True
        )
        self.nodes_lstm = tf.keras.layers.LSTM(
            units, name="nodes_lstm", return_sequences=True, time_major=False
        )
        self.embedding = tf.keras.layers.Dense(units=self.units, name="embedding")
        self.build_feature_columns()
        self.attention = BahdanauAttention(units, name="attention")
        self.dense_h = tf.keras.layers.Dense(units, name=f"lstm_initial_h")