docs: public SDVariable batchNorm(SDVariable input, SDVariable mean,
                                  SDVariable variance, SDVariable gamma,
                                  SDVariable beta,
                                  boolean applyGamma, boolean applyBeta,
                                  double epsilon, int... axis) {
          BatchNorm batchNorm = BatchNorm.builder()
                  .inputFunctions(new SDVariable[]{input, mean, variance, gamma, beta})
                  .applyGamma(applyGamma)
                  .applyBeta(applyBeta)
                  .epsilon(epsilon)
                  .sameDiff(sameDiff())
                  .axis(axis)
                  .build();
  
          val outputVars = batchNorm.outputVariables();
          return outputVars[0];
      }