red-chainer 0.3.2 → 0.4.0

data/lib/chainer/functions/math/basic_math.rb

@@ -1,66 +1,71 @@
 module Chainer
   module Functions
     module Math
- 
-      class Neg < ::Chainer::Function
+      class Neg < ::Chainer::FunctionNode
+        def label
+          '__neg__'
+        end
+
         def forward(x)
-          retain_inputs([])
           [Utils::Array.force_array(-x[0])]
         end
 
-        def backward(x, gy)
-          [Utils::Array.force_array(-gy[0])]
+        def backward(indexes, gy)
+          [-gy[0]]
         end
       end
 
-      class Add < ::Chainer::Function
+      class Add < ::Chainer::FunctionNode
         def forward(x)
-          retain_inputs([])
           [Utils::Array.force_array(x[0] + x[1])]
         end
 
-        def backward(x, gy)
+        def backward(indexes, gy)
           [gy[0], gy[0]]
         end
       end
 
-      class AddConstant < ::Chainer::Function
+      class AddConstant < ::Chainer::FunctionNode
         def initialize(value)
           @value = value
         end
 
         def forward(x)
-          retain_inputs([])
           [Utils::Array.force_array(x[0] + @value)]
         end
 
-        def backward(x, gy)
+        def backward(indexes, gy)
           [gy[0]]
         end
       end
- 
-      class Sub < ::Chainer::Function
+
+      class Sub < ::Chainer::FunctionNode
+        def label
+          '_ - _'
+        end
+
         def forward(x)
-          retain_inputs([])
           [Utils::Array.force_array(x[0] - x[1])]
         end
 
-        def backward(x, gy)
-          [gy[0], Utils::Array.force_array(-gy[0])]
+        def backward(indexes, gy)
+          [gy[0], -gy[0]]
         end
       end
 
-      class Mul < ::Chainer::Function
+      class Mul < ::Chainer::FunctionNode
         def forward(x)
+          retain_inputs([0, 1])
           [Utils::Array.force_array(x[0] * x[1])]
         end
 
-        def backward(x, gy)
-          [Utils::Array.force_array(gy[0] * x[1]), Utils::Array.force_array(gy[0] * x[0])]
+        def backward(indexes, gy)
+          xs = get_retained_inputs
+          indexes.map { |i| gy[0] * xs[1 - i] }
         end
       end
 
-      class MulConstant < ::Chainer::Function
+      class MulConstant < ::Chainer::FunctionNode
         def initialize(value)
           @value = value
         end
@@ -69,23 +74,23 @@ module Chainer
           [Utils::Array.force_array(@value * x[0])]
         end
 
-        def backward(x, gy)
-          [Utils::Array.force_array(@value * gy[0])]
+        def backward(indexes, gy)
+          [gy[0] * @value]
         end
       end
-      
-      class Div < ::Chainer::Function
+
+      class Div < ::Chainer::FunctionNode
         def forward(x)
           [Utils::Array.force_array(x[0] / x[1])]
         end
 
-        def backward(x, gy)
+        def backward(indexes, gy)
           gx0 = Utils::Array.force_array(gy[0] / x[1])
           [gx0, Utils::Array.force_array(-1 * gx0 * x[0] / x[1])]
         end
       end
-      
-      class PowVarVar < ::Chainer::Function
+
+      class PowVarVar < ::Chainer::FunctionNode
         def forward(x)
           @y = Utils::Array.force_array(x[0] ** x[1])
           [@y]
@@ -94,12 +99,13 @@ module Chainer
         def backward(x, gy)
           one = x[1].class.ones[0]
           gx0 = Utils::Array.force_array(x[1] * (x[0] ** (x[1] - one)) * gy[0])
-          gx1 = Utils::Array.force_array(Numo::NMath.log(x[0]) * @y * gy[0])
+          xm = Chainer.get_array_module(x[0])
+          gx1 = Utils::Array.force_array(xm::NMath.log(x[0]) * @y * gy[0])
           [gx0, gx1]
         end
       end
 
-      class PowVarConst < ::Chainer::Function
+      class PowVarConst < ::Chainer::FunctionNode
         def initialize(value)
           @value = value
         end
@@ -113,7 +119,7 @@ module Chainer
           gx = @value * (x[0] ** val_1) * gy[0]
           [Utils::Array.force_array(gx)]
         end
-      end 
+      end
     end
   end
 end

data/lib/chainer/functions/math/exp.rb

@@ -0,0 +1,28 @@
+module Chainer
+  module Functions
+    module Math
+      class Exp < Chainer::FunctionNode
+        # Elementwise exponential function.
+        def self.exp(x)
+          self.new.apply([x]).first
+        end
+
+        def label
+          'exp'
+        end
+
+        def forward(x)
+          retain_inputs([])
+          retain_outputs([0])
+          xm = Chainer.get_array_module(x.first)
+          [Utils::Array.force_array(xm::NMath.exp(x.first))]
+        end
+
+        def backward(indexes, gy)
+          y = get_retained_outputs.first
+          [y * gy.first]
+        end
+      end
+    end
+  end
+end

data/lib/chainer/functions/math/identity.rb

@@ -2,7 +2,7 @@ module Chainer
   module Functions
     module Math
       # Identity function.
-      class Identity < Chainer::Function
+      class Identity < Chainer::FunctionNode
         def check_type_forward(in_types)
           # pass
         end
@@ -12,13 +12,14 @@ module Chainer
           return xs
         end
 
-        def backward(xs, gys)
+        def backward(indexes, gys)
           return gys
         end
 
         # Just returns input variables.
         def self.identity(*inputs)
-          self.new.(*inputs)
+          ret = self.new.apply(inputs)
+          ret.size == 1 ? ret[0] : ret
         end
       end
     end

data/lib/chainer/functions/math/sum.rb

@@ -0,0 +1,52 @@
+module Chainer
+  module Functions
+    module Math
+      # Sum of array elements over a given axis.
+      class Sum < Chainer::FunctionNode
+        # Sum of array elements over a given axis
+        #
+        # @param [Chainer::Variable] x Elements to sum
+        # @param [nil, Integer, Array<Integer>] axis Axis which a sum is performed
+        # @param[boolean] keepdims If `true`, the specified axes are remained as axes of length one
+        # @return [Chainer::Variable] Output variable
+        def self.sum(x, axis: nil, keepdims: false)
+          Sum.new(axis: axis, keepdims: keepdims).apply([x]).first
+        end
+
+        def initialize(axis: nil, keepdims: false)
+          if axis.nil?
+            @axis = nil
+          elsif axis.is_a?(Integer)
+            @axis = [axis]
+          elsif axis.is_a?(::Array) && axis.all? { |e| e.is_a?(Integer) }
+            raise ArgumentError, "duplicate value in axis: #{axis}" unless axis.uniq.size == axis.size
+            @axis = axis
+          else
+            raise TypeError, 'nil, Integer or Array of int are required'
+          end
+
+          @keepdims = keepdims
+        end
+
+        def forward(inputs)
+          x = inputs.first
+          ret = x.sum(axis: @axis, keepdims: @keepdims)
+          ret = x.class.cast(ret)
+          [ret]
+        end
+
+        def backward(indexes, grad_outputs)
+          gy = grad_outputs.first
+          ndim = @inputs.first.shape.size
+          unless ndim == 0 || @axis.nil? || @keepdims
+            actual_axis = @axis.map { |axis| axis >= 0 ? axis : axis + ndim  }
+            shape = gy.shape
+            actual_axis.sort.each { |axis| shape.insert(axis, 1) }
+            gy = Chainer::Functions::Array::Reshape.reshape(gy, shape)
+          end
+          [Chainer::Functions::Array::BroadcastTo.broadcast_to(gy, @inputs.first.shape)]
+        end
+      end
+    end
+  end
+end

data/lib/chainer/functions/noise/dropout.rb

@@ -1,7 +1,8 @@
 module Chainer
   module Functions
     module Noise
-      class Dropout < Chainer::Function
+      class Dropout < Chainer::FunctionNode
+        attr_reader :mask
         # Drops elements of input variable randomly.
         #
         # This function drops input elements randomly with probability `ratio` and
@@ -12,7 +13,7 @@ module Chainer
         # @param [float] ratio Dropout ratio. The ``ratio`` must be `0.0 <= ratio < 1.0`.
         # @return [Chainer::Variable] Output variable.
         def self.dropout(x, ratio: 0.5)
-          Chainer.configuration.train ? self.new(ratio).(x) : x
+          Chainer.configuration.train ? self.new(ratio).apply([x])[0] : Chainer::Variable.as_variable(x)
         end
 
         def initialize(dropout_ratio)
@@ -23,7 +24,6 @@ module Chainer
         end
 
         def forward(x)
-          retain_inputs([])
           unless self.instance_variable_defined?(:@mask)
             scale = x[0].class[*[1.0 / (1 - @dropout_ratio)]][0]
             flag = x[0].class.new(*x[0].shape).rand >= @dropout_ratio
@@ -36,7 +36,23 @@ module Chainer
         end
 
         def backward(x, gy)
-          [gy[0] * @mask]
+          DropoutGrad.new(@mask).apply(gy)
+        end
+      end
+
+      # Computes the gradient of the Dropout function.
+      class DropoutGrad < Chainer::FunctionNode
+        def initialize(mask)
+          @mask = mask
+        end
+
+        def forward(inputs)
+          y = inputs.first * @mask
+          [y]
+        end
+
+        def backward(indexes, gy)
+          DropoutGrad.new(@mask).apply(gy)
         end
       end
     end

data/lib/chainer/functions/normalization/batch_normalization.rb

@@ -1,134 +1,287 @@
 module Chainer
   module Functions
     module Normalization
-      class BatchNormalizationFunction < Chainer::Function
+      module Calculation
+        def apply_bn_fwd(xp, x, mean, inv_std, gamma, beta)
+          # NOTE: all arguments should be broadcasted to x.shape
+          # (mean, inv_std, gamma, and beta have to already be expanded)
+          x_hat = x_hat(x, mean, inv_std)
+          y = gamma * x_hat
+          y += beta
+          y
+        end
+
+        def x_hat(x, mean, inv_std)
+          x_mu = x - mean
+          x_mu *= inv_std
+          x_mu
+        end
+
+        def zero_if_none(xp, x, shape, dtype)
+          # TODO: Return broadcasted 0 instead of a zeroed array.
+          x.nil? ? dtype.zeros(*shape) : x
+        end
+      end
+
+      class BatchNormalization < Chainer::FunctionNode
+        include Calculation
         attr_reader :running_mean, :running_var
-        # Batch normalization function with fixed statistics.
-        # This is a variant of batch normalization, where the mean and variance
-        # statistics are given by the caller as fixed variables. This is
-        # used on testing mode of the batch normalization layer, where batch
-        # statistics cannot be used for prediction consistency.
-        #
-        # @param [Chainer::Variable] x Input variable.
-        # @param [Chainer::Variable] gamma Scaling parameter of normalized data.
-        # @param [Chainer::Variable] beta Shifting parameter of scaled normalized data.
-        # @param [Chainer::Variable] mean Shifting parameter of input.
-        # @param [Chainer::Variable] var Square of scaling parameter of input.
-        # @param [float] eps Epsilon value for numerical stability.
-        def self.fixed_batch_normalization(x, gamma, beta, mean, var, eps: 2e-5)
-          old_train = Chainer.configuration.train
-          Chainer.configuration.train = false
-          norm = self.new(eps: eps, mean: nil, var: nil, decay: 0.0).(x, gamma, beta, mean, var)
-          Chainer.configuration.train = old_train
-          norm
-        end
-      
-        def initialize(eps: 2e-5, mean: nil, var: nil, decay: 0.9) 
+
+        def self.batch_normalization(x, gamma, beta, eps: 2e-5, running_mean: nil, running_var: nil, decay: 0.9)
+          BatchNormalization.new(eps: eps, mean: running_mean, var: running_var, decay: decay).apply([x, gamma, beta])[0]
+        end
+
+        def initialize(eps: 2e-5, mean: nil, var: nil, decay: 0.9)
+          @mean = nil
+          @inv_std = nil
+
           @running_mean = mean
           @running_var = var
           @eps = eps
-          @mean_cache = nil
           @decay = decay
         end
 
         def forward(inputs)
-          x, gamma, beta = inputs[0], inputs[1], inputs[2]
-          if Chainer.configuration.train
-            if @running_mean.nil?
-              @running_mean = Numo::NArray[*gamma].new_zeros
-              @running_var = Numo::NArray[*gamma].new_zeros
-            else
-              @running_mean = Numo::NArray[*@running_mean]
-              @running_var = Numo::NArray[*@running_var]
-            end
-          elsif inputs.size == 5
-            @fixed_mean = inputs[3]
-            @fixed_var = inputs[4]
+          retain_inputs([0, 1])
+          x, gamma, beta = inputs
+          xp = Chainer.get_array_module(x)
+
+          if @running_mean.nil?
+            @running_mean = xp::NArray[*gamma].new_zeros
+            @running_var = xp::NArray[*gamma].new_zeros
           end
 
+          # expander inserts singleton dimensions to gamma and beta so that they
+          # can be broadcasted with x.
           head_ndim = gamma.ndim + 1
-          gamma_expander = [1] + gamma.shape + [1] * (x.ndim - head_ndim)
-          gamma = gamma.reshape(*gamma_expander)
-          beta_expander = [1] + beta.shape + [1] * (x.ndim - head_ndim)
-          beta = beta.reshape(*beta_expander)
-        
-          if Chainer.configuration.train
-            axis = [0] + (head_ndim...(x.ndim)).to_a
-            mean = x.mean(axis: axis)
-            # FIXME: numpy.var
-            var = x.var(axis: axis)
-            var += @eps
-          else
-            mean = @fixed_mean
-            var = @fixed_var + @eps
+          # TODO: expander = (None, Ellipsis) + (None,) * (x.ndim - head_ndim)
+          suffix = [1] * (x.ndim - head_ndim)
+          expander = -> (arr) do
+            shape = [1] + arr.shape + suffix
+            arr.reshape(*shape)
           end
+          @expander = expander
+          @axis = [0] + (head_ndim...(x.ndim)).to_a
 
-          @std = Numo::NMath.sqrt(var)
+          gamma = expander.(gamma)
+          beta = expander.(beta)
+          @mean = x.mean(axis: @axis)
 
-          mean_expander = [1] + mean.shape + [1] * (x.ndim - head_ndim)
-          x_mu = x - mean.reshape(*mean_expander)
-          std_expander = [1] + @std.shape + [1] * (x.ndim - head_ndim)
-          x_mu /= @std.reshape(*std_expander)
-          @x_hat = x_mu
-          y = gamma * @x_hat
-          y += beta
+          # TODO: Numo::Array can not be specified standard deviation
+          var = ((x - x.mean(axis: @axis, keepdims: true)) ** 2).mean(axis: @axis)
 
-          if Chainer.configuration.train
-            m = x.size.div(gamma.size)
-            adjust = m / [m - 1.0, 1.0].max
-            @running_mean *= @decay
-            temp_ar = Numo::NArray[*mean]
-            temp_ar *= (1 - @decay)
-            @running_mean += temp_ar
-            
-            @running_var *= @decay
-            temp_ar = Numo::NArray[*var]
-            temp_ar *= ((1 - @decay) * adjust)
-            @running_var += temp_ar
-          end
+          var += @eps
+          @inv_std = var ** (-0.5)
+
+          y = apply_bn_fwd(xp, x, expander.(@mean), expander.(@inv_std), gamma, beta)
+          # Update running statistics
+          m = x.size.div(gamma.size)
+          adjust = m / [m - 1.0, 1.0].max
+          @running_mean *= @decay
+          @running_mean += (1 - @decay) * @mean
+          @running_var *= @decay
+          @running_var += (1 - @decay) * adjust * var
 
-          [y,]
+          [y]
+        end
+
+        def backward(indexes, grad_outputs)
+          x, gamma = get_retained_inputs
+          gy, = grad_outputs
+
+          # hatappi debug
+          #@mean = @mean.class.new(@mean.shape).seq
+          #@inv_std = @inv_std.class.new(@inv_std.shape).seq
+          #x.data = x.data.class.new(x.shape).seq
+          #gamma.data = gamma.data.class.new(gamma.shape).seq
+          #gy.data = gy.data.class.new(gy.shape).seq
+
+          f = BatchNormalizationGrad.new(@eps, @expander, @axis, @mean, @inv_std)
+          f.(x, gamma, gy)
+        end
+      end
+
+      class BatchNormalizationGrad < Function
+        include Calculation
+
+        def initialize(eps, expander, axis, mean, inv_std)
+          @eps = eps
+          @expander = expander
+          @axis = axis
+          @mean = mean
+          @inv_std = inv_std
+        end
+
+        def forward(inputs)
+          retain_inputs([0, 1, 2])
+          x, gamma, gy = inputs
+          expander = @expander
+
+          inv_m = gamma.class.new.fill(1.0 / x.size.div(gamma.size))
+          xp = Chainer.get_array_module(x)
+
+          gbeta = gy.sum(axis: @axis)
+          x_hat = x_hat(x, expander.(@mean), expander.(@inv_std))
+          ggamma = (gy * x_hat).sum(axis: @axis)
+          gx = expander.(gamma * @inv_std) * (gy - (x_hat * expander.(ggamma) + expander.(gbeta)) * inv_m)
+
+          retain_outputs([0, 1])
+          [gx, ggamma, gbeta]
         end
 
         def backward(inputs, grad_outputs)
-          x, gamma = inputs[0], inputs[1]
-          gy = grad_outputs[0]
+          expander = @expander
+
+          x, gamma, gy = inputs
+          gx1, ggamma1, = output_data
+          ggx1, gggamma1, ggbeta1 = grad_outputs
+          xp = Chainer.get_array_module(x)
+
+          # auxiliary values
+          inv_m = gamma.class.new.fill(1.0 / x.size.div(gamma.size))
+          r = ggx1.nil? ? 0 : (gx1 * ggx1).sum(axis: @axis)
+          coeff = gamma * @inv_std
+          coeff_m = coeff * inv_m
+          x_hat = x_hat(x, expander.(@mean), expander.(@inv_std))
+
+          # handle None in output gradients
+          ggx1 = zero_if_none(xp, ggx1, x.shape, x.class)
+          gggamma1 = zero_if_none(xp, gggamma1, gamma.shape, gamma.class)
+          ggbeta1 = zero_if_none(xp, ggbeta1, gamma.shape, gamma.class)
+
+          gggamma2 = gggamma1 - coeff_m * (x_hat * ggx1).sum(axis: @axis)
+        	ggbeta2 = ggbeta1 - coeff_m * ggx1.sum(axis: @axis)
+
+          ggamma2 = r / gamma
+
+          gx_hat2 = (expander.(gggamma2) * gy - expander.(coeff_m * ggamma1) * ggx1)
+          gstd2 = -@inv_std * (r + (x_hat * gx_hat2).sum(axis: @axis))
+          gmean2 = -@inv_std * gx_hat2.sum(axis: @axis)
+          gx2 = expander.(@inv_std) * gx_hat2 + inv_m * (expander.(gmean2) + x_hat * expander.(gstd2))
+          ggy2 = (expander.(gggamma2) * x_hat + expander.(ggbeta2) + expander.(coeff) * ggx1)
+
+          [gx2, ggamma2, ggy2]
+        end
+      end
+
+      class FixedBatchNormalization < FunctionNode
+        include Calculation
+
+        attr_reader :inv_var
+
+        def self.fixed_batch_normalization(x, gamma, beta, mean, var, eps: 2e-5)
+          FixedBatchNormalization.new(eps: eps).apply([x, gamma, beta, mean, var]).first
+        end
+
+        def initialize(eps: 2e-5)
+          @inv_std = nil
+          @inv_var = nil
+          @eps = eps
+        end
+
+        def forward(inputs)
+          retain_inputs([0, 1, 3, 4])
+          x, gamma, beta, mean, var = inputs
+          xp = Chainer.get_array_module(x)
+
+          # expander inserts singleton dimensions to gamma and beta so that they
+          # can be broadcasted with x.
           head_ndim = gamma.ndim + 1
-          m = gamma.class[x.size.div(gamma.size)][0]
-          axis = [0] + (head_ndim...(x.ndim)).to_a
-
-          if inputs.size == 5
-            mean = inputs[3]
-            var = inputs[4]
-            std = Numo::NMath.sqrt(var)
-            gs = gamma / std
-            gbeta = gy.sum(axis: axis)
-
-            mean_expander = [1] + mean.shape + [1] * (x.ndim - head_ndim)
-            x_mu = x - mean.reshape(*mean_expander)
-            std_expander = [1] + std.shape + [1] * (x.ndim - head_ndim)
-            x_mu /= std.reshape(*std_expander)
-            x_hat = x_mu
-            ggamma = (gy * x_hat).sum(axis: axis)
-            gmean = -gs * gbeta
-            gvar = -0.5 * gamma / var * ggamma
-            gs_expander = [1] + gs.shape + [1] * (x.ndim - head_ndim)
-            gx = gs.reshape(*gs_expander)
-            return [gx, ggamma, gbeta, gmean, gvar]
+          # TODO: expander = (None, Ellipsis) + (None,) * (x.ndim - head_ndim)
+          suffix = [1] * (x.ndim - head_ndim)
+          expander = -> (arr) do
+            shape = [1] + arr.shape + suffix
+            arr.reshape(*shape)
           end
+          @expander = expander
+          @axis = [0] + (head_ndim...(x.ndim)).to_a
+
+          gamma = expander.(gamma)
+          beta = expander.(beta)
+          var += @eps
+          @inv_var = var.reciprocal
+          @inv_std = xp::NMath.sqrt(@inv_var)
 
-          gbeta = gy.sum(axis: axis)
-          ggamma = (gy * @x_hat).sum(axis: axis)
-          tmp = (gamma / @std)
-          tmp_expander = [1] + tmp.shape + [1] * (x.ndim - head_ndim)
-          tmp = tmp.reshape(*tmp_expander)
+          y = apply_bn_fwd(xp, x, expander.(mean), expander.(@inv_std), gamma, beta)
+          [y]
+        end
 
-          ggamma_expander = [1] + ggamma.shape + [1] * (x.ndim - head_ndim)
-          gbeta_expander = [1] + gbeta.shape + [1] * (x.ndim - head_ndim)
-          
-          gx = tmp * (gy - (@x_hat * ggamma.reshape(*ggamma_expander) + gbeta.reshape(*gbeta_expander)) / m )
+        def backward(indexes, grad_outputs)
+          x, gamma, mean, var = get_retained_inputs
+          gy, = grad_outputs
+          f = FixedBatchNormalizationGrad.new(@eps, @expander, @axis, @inv_std, @inv_var)
+          f.(x, gamma, mean, var, gy)
+        end
+      end
 
-          [gx, ggamma, gbeta]
+      class FixedBatchNormalizationGrad < Function
+        include Calculation
+
+        def initialize(eps, expander, axis, inv_std, inv_var)
+          @eps = eps
+          @expander = expander
+          @axis = axis
+          @inv_std = inv_std
+          @inv_var = inv_var
+        end
+
+        def forward(inputs)
+          retain_inputs([0, 1, 2, 4])
+          x, gamma, mean, var, gy = inputs
+          expander = @expander
+          xp = Chainer.get_array_module(x)
+
+          if @inv_std.nil? || @inv_var.nil?
+            @inv_var = (var + @eps).reciprocal
+            @inv_std = xp::NMath.sqrt(@inv_var)
+          end
+
+          @gamma_over_std = gamma * @inv_std
+          x_hat = x_hat(x, expander.(mean), expander.(@inv_std))
+
+          gx = expander.(@gamma_over_std) * gy
+          gbeta = gy.sum(axis: @axis)
+          ggamma = (x_hat * gy).sum(axis: @axis)
+          gmean = -@gamma_over_std * gbeta
+          gvar = -0.5 * gamma * @inv_var * ggamma
+
+          retain_outputs([0, 1, 2, 3, 4])
+          [gx, ggamma, gbeta, gmean, gvar]
+        end
+
+        def backward(inputs, grad_outputs)
+          x, gamma, mean, _, gy = inputs
+          ggx1, gggamma1, ggbeta1, ggmean1, ggvar1 = grad_outputs
+          gx1, ggamma1, gbeta1, gmean1, gvar1 = output_data
+
+          # Handle None in output gradients.
+          xp = Chainer.get_array_module(x)
+          ggx1 = zero_if_none(xp, ggx1, x.shape, x.class)
+          gggamma1 = zero_if_none(xp, gggamma1, gamma.shape, gamma.class)
+          ggbeta1 = zero_if_none(xp, ggbeta1, gamma.shape, gamma.class)
+          ggmean1 = zero_if_none(xp, ggmean1, mean.shape, mean.class)
+          ggvar1 = zero_if_none(xp, ggvar1, mean.shape, mean.class)
+
+          expander = @expander
+          x_hat = x_hat(x, expander.(mean), expander.(@inv_std))
+          tmp = -0.5 * ggvar1
+
+          gamma_over_var = gamma * @inv_var
+          g_gamma_over_var = tmp * ggamma1
+
+          gggamma2 = gggamma1 + tmp * gamma_over_var
+          gx_hat = gy * expander.(gggamma2)
+          gx2 = expander.(@inv_std) * gx_hat
+          gmean2 = -@inv_std * gx_hat.sum(axis: @axis)
+
+          g_gamma_over_std = (ggx1 * gy).sum(axis: @axis) - ggmean1 * gbeta1
+          ggbeta2 = ggbeta1 - ggmean1 * @gamma_over_std
+          ggy2 = (expander.(gggamma2) * x_hat + expander.(ggbeta2) + expander.(@gamma_over_std) * ggx1)
+
+          ggamma2 = (@inv_var * g_gamma_over_var + @inv_std * g_gamma_over_std)
+          gvar2 = -(ggamma2 * gamma_over_var + 0.5 * @inv_var * ((x_hat * gx_hat).sum(axis: @axis) - @gamma_over_std * g_gamma_over_std))
+
+          [gx2, ggamma2, gmean2, gvar2, ggy2]
         end
       end
     end