Function BwTensorDescInferFn

oneflow/user/ops/normalization_op.cpp:433–458

namespace {

Maybe<void> BwTensorDescInferFn(user_op::InferContext* ctx) {
#ifdef WITH_CUDA
  // assume cudnn is enabled
  CHECK_GE_OR_RETURN(ctx->Attr<float>("epsilon"), CUDNN_BN_MIN_EPSILON);
#endif
  const user_op::TensorDesc& x = ctx->InputTensorDesc("x", 0);
  const Shape& x_shape = x.shape();
  const user_op::TensorDesc& dy = ctx->InputTensorDesc("dy", 0);
  CHECK_EQ_OR_RETURN(dy.shape(), x_shape);
  if (ctx->has_input("y", 0)) {
    const user_op::TensorDesc& y = ctx->InputTensorDesc("y", 0);
    CHECK_EQ_OR_RETURN(y.shape(), x_shape);
  }
  *ctx->MutOutputTensorDesc("dx", 0) = x;
  if (ctx->has_output("addend_diff", 0)) { *ctx->MutOutputTensorDesc("addend_diff", 0) = x; }
  const Shape param_shape({x_shape.At(ctx->Attr<int32_t>("axis"))});
  const auto CheckParamTensorDesc = MakeCheckParamTensorDescFn(ctx, param_shape);
  const auto SetParamTensorDesc = MakeSetParamTensorDescFn(ctx, param_shape);
  JUST(CheckParamTensorDesc("mean"));
  JUST(CheckParamTensorDesc("inv_variance"));
  JUST(CheckParamTensorDesc("gamma"));
  JUST(CheckParamTensorDesc("beta"));
  JUST(SetParamTensorDesc("gamma_diff"));
  JUST(SetParamTensorDesc("beta_diff"));
  return Maybe<void>::Ok();
}

Maybe<void> BwDataTypeInferFn(user_op::InferContext* ctx) {
  const user_op::TensorDesc& x = ctx->InputTensorDesc("x", 0);