class LayerNormStackSmallOp(HybridBlock):
    """Applies layer normalization to the n-dimensional input array.
    Stack bcast/reduce
    """
    def __init__(self, axis=1, epsilon=1e-5, center=True, scale=True,
                 beta_initializer='zeros', gamma_initializer='ones',
                 in_channels=0, prefix=None, params=None):
        super(LayerNormStackSmallOp, self).__init__(prefix=prefix, params=params)
        self._kwargs = {'eps': epsilon, 'axis': axis}
        self._axis = axis
        self._epsilon = epsilon
        self._center = center
        self._scale = scale
        assert in_channels != 0, "in_channels == 0 is currently not supported"
        if self._center:
            self.gamma = self.params.get('gamma', grad_req='write' if scale else 'null',
                                         shape=(in_channels,), init=gamma_initializer,
                                         allow_deferred_init=True)
        if self._scale:
            self.beta = self.params.get('beta', grad_req='write' if center else 'null',
                                        shape=(in_channels,), init=beta_initializer,
                                        allow_deferred_init=True)

    def moments(self, F, data):
        mean = F.mean(data=data, axis=self._axis, keepdims=True)
        var = F.mean(F.square(F.broadcast_minus(data, mean)),
                     axis=self._axis, keepdims=True)
        return mean, var

    def hybrid_forward(self, F, data, gamma, beta):
        if not self._center and not self._scale:
            return data
        mean, var = self.moments(F, data)
        norm_data = F.broadcast_minus(data, mean)
        norm_data = F.broadcast_mul(norm_data, mx.sym.rsqrt(var + self._epsilon))
        norm_data = F.broadcast_mul(norm_data, gamma)
        norm_data = F.broadcast_add(norm_data, beta)
        return norm_data