src/transformers/models/gptj/modeling_tf_gptj.py/TFGPTJBlock/__init__
class TFGPTJBlock:    def __init__(self, config: GPTJConfig, **kwargs):
        super().__init__(**kwargs)
        inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd
        self.ln_1 = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="ln_1")
        self.attn = TFGPTJAttention(config, name="attn")
        self.mlp = TFGPTJMLP(inner_dim, config, name="mlp")