import numpy as np
import tensorflow as tf
H = 2
N = 2
M = 3
BS = 10

def my_softmax(arr):
    max_elements = np.reshape(np.max(arr, axis = 2), (BS, N, 1))
    arr = arr - max_elements
    exp_array = np.exp(arr)
    print (exp_array)
    sum_array = np.reshape(
np.sum(exp_array, axis=2)