reserve_spaces[1] = op.input(4);  // pop_var
    if (data_format == "NCHW") {
      x = Transpose(scope, x, {0, 2, 3, 1});
      grad_y = Transpose(scope, grad_y, {0, 2, 3, 1});
    } else if (data_format == "NCDHW") {
      x = Transpose(scope, x, {0, 2, 3, 4, 1});
      grad_y = Transpose(scope, grad_y, {0, 2, 3, 4, 1});
    }

    StringPiece target_data_format;
    if (data_format == "NCHW" || data_format == "NHWC") {

  auto var = ops::OnesLike(scope_, scale);

  if (!channel_first) {
    attrs.data_format_ = (shape.dims() == 5) ? "NDHWC" : "NHWC";
  } else {
    attrs.data_format_ = (shape.dims() == 5) ? "NCDHW" : "NCHW";
  }

  auto y = FusedBatchNormV3(scope_, x, scale, offset, mean, var, attrs);
  RunTest({x, scale, offset}, {shape, scale_shape, scale_shape}, {y.y},
          {shape});
}

    ConfinedAttr<I64ArrayAttr, [ArrayMinCount<5>]>:$ksize,
    ConfinedAttr<I64ArrayAttr, [ArrayMinCount<5>]>:$strides,
    TF_AnyStrAttrOf<["SAME", "VALID"]>:$padding,
    DefaultValuedOptionalAttr<TF_AnyStrAttrOf<["NDHWC", "NCDHW"]>, "\"NDHWC\"">:$data_format
  );

  let results = (outs
    Res<TF_FloatTensor, [{The average pooled output tensor.}]>:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

  bool is_pool2d = rank == 4;
  if (verify_batch_channel_dims(0, rank - 1)) {
    *data_format = is_pool2d ? "NHWC" : "NDHWC";
  } else if (verify_batch_channel_dims(0, 1)) {
    *data_format = is_pool2d ? "NCHW" : "NCDHW";
  } else {
    return false;
  }

  if (rw.getWindowDilations().has_value() &&
      !(rw.getWindowDilations()->isSplat() &&
        rw.getWindowDilations()->getSplatValue<APInt>() == 1))

// CHECK-SAME:                                %[[VAL_0:.*]]: tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32> {
// CHECK:           %[[VAL_1:.*]] = "tf.AvgPool3D"(%[[VAL_0]]) <{data_format = "NCDHW", ksize = [1, 1, 3, 3, 3], padding = "VALID", strides = [1, 1, 2, 2, 2]}> : (tensor<4x3x16x16x16xf32>) -> tensor<4x3x7x7x7xf32>
// CHECK:           return %[[VAL_1]] : tensor<4x3x7x7x7xf32>
// CHECK:         }

// CHECK-SAME:        -> tensor<2x7x4x3x5xf16>
// CHECK:           return [[CONV16]]
func.func @avgpool_3d_ncdhw_format(%arg0: tensor<2x7x4x12x21xf16>) -> tensor<2x7x4x3x5xf16> {
  %0 = "tf.AvgPool3D"(%arg0) {data_format = "NCDHW", ksize = [1, 1, 1, 2, 2], padding = "VALID", strides = [1, 1, 1, 4, 4]} : (tensor<2x7x4x12x21xf16>) -> tensor<2x7x4x3x5xf16>
  func.return %0 : tensor<2x7x4x3x5xf16>
}

// -----