Your J_style_layer code seems correct. How you are reshaping the tensors from (1, n_H, n_W, n_C) to (n_C, n_H * n_W)? You have to use both, tf.reshape and tf.transpose. First, use tf.reshape and then tf.transpose. Maybe this thread seems relevant to your query.
2 Likes