Source:
public Pair<Gradient, INDArray> backpropGradient(INDArray epsilon, LayerWorkspaceMgr workspaceMgr) { assertInputSet(true); Pair<INDArray, INDArray> zAndPreNorm = preOutputWithPreNorm(true, workspaceMgr); INDArray z = zAndPreNorm.getFirst(); INDArray preNorm = zAndPreNorm.getSecond(); INDArray delta = layerConf().getActivationFn().backprop(z, epsilon).getFirst(); if (maskArray != null) { applyMask(delta); } Gradient ret = new DefaultGradient(); if(hasBias()){ INDArray biasGrad = gradientViews.get(DefaultParamInitializer.BIAS_KEY); delta.sum(biasGrad, 0); ret.gradientForVariable().put(DefaultParamInitializer.BIAS_KEY, biasGrad); } INDArray W = getParamWithNoise(DefaultParamInitializer.WEIGHT_KEY, true, workspaceMgr); INDArray epsilonNext = workspaceMgr.createUninitialized(ArrayType.ACTIVATION_GRAD, new long[]{W.size(0), delta.size(0)}, 'f'); if(hasLayerNorm()) { INDArray g = getParamWithNoise(DefaultParamInitializer.GAIN_KEY, true, workspaceMgr); INDArray dldg = gradientViews.get(DefaultParamInitializer.GAIN_KEY); Nd4j.getExecutioner().exec(new LayerNormBp(preNorm, g, delta, delta, dldg, 1)); ret.gradientForVariable().put(DefaultParamInitializer.GAIN_KEY, dldg); } epsilonNext = W.mmuli(delta.transpose(),epsilonNext).transpose(); INDArray weightGrad = gradientViews.get(DefaultParamInitializer.WEIGHT_KEY); Nd4j.gemm(input.castTo(weightGrad.dataType()), delta, weightGrad, true, false, 1.0, 0.0); ret.gradientForVariable().put(DefaultParamInitializer.WEIGHT_KEY, weightGrad); weightNoiseParams.clear(); epsilonNext = backpropDropOutIfPresent(epsilonNext); return new Pair<>(ret, epsilonNext); }

Target:
public Pair<Gradient, INDArray> backpropGradient(INDArray epsilon, LayerWorkspaceMgr workspaceMgr) { assertInputSet(true); Pair<INDArray, INDArray> zAndPreNorm = preOutputWithPreNorm(true, true, workspaceMgr); INDArray z = zAndPreNorm.getFirst(); INDArray preNorm = zAndPreNorm.getSecond(); INDArray delta = layerConf().getActivationFn().backprop(z, epsilon).getFirst(); if (maskArray != null) { applyMask(delta); } Gradient ret = new DefaultGradient(); if(hasBias()){ INDArray biasGrad = gradientViews.get(DefaultParamInitializer.BIAS_KEY); delta.sum(biasGrad, 0); ret.gradientForVariable().put(DefaultParamInitializer.BIAS_KEY, biasGrad); } INDArray W = getParamWithNoise(DefaultParamInitializer.WEIGHT_KEY, true, workspaceMgr); INDArray epsilonNext = workspaceMgr.createUninitialized(ArrayType.ACTIVATION_GRAD, new long[]{W.size(0), delta.size(0)}, 'f'); if(hasLayerNorm()) { INDArray g = getParam(DefaultParamInitializer.GAIN_KEY); INDArray dldg = gradientViews.get(DefaultParamInitializer.GAIN_KEY); Nd4j.getExecutioner().exec(new LayerNormBp(preNorm, g, delta, delta, dldg, 1)); ret.gradientForVariable().put(DefaultParamInitializer.GAIN_KEY, dldg); } epsilonNext = W.mmuli(delta.transpose(),epsilonNext).transpose(); INDArray weightGrad = gradientViews.get(DefaultParamInitializer.WEIGHT_KEY); Nd4j.gemm(input.castTo(weightGrad.dataType()), delta, weightGrad, true, false, 1.0, 0.0); ret.gradientForVariable().put(DefaultParamInitializer.WEIGHT_KEY, weightGrad); weightNoiseParams.clear(); epsilonNext = backpropDropOutIfPresent(epsilonNext); return new Pair<>(ret, epsilonNext); }