module {
  // For XLA weight-only per-channel depthwise convolution, tensor shape should have
  // transformed from [H,W,C,M] to [H,W,1,CxM],
  func.func @depthwise_conv(%arg0: tensor<1x3x4x3xf32>) -> (tensor<*xf32>) {

// Processes tensors with NCHW format (== (batch, channel, height, weight)) by
// converting them to NHWC formats along with extra optimizations such as
// constant folding the transpose->convolution pattern. This is useful when
// downstream pipeline (e.g. XLA) is more optimized when accepting NHWC formats.
void AddProcessNchwTensorPasses(OpPassManager& pm);

    }

    // If all quantized or floating point then types are consistent.
    // Int is valid in combination with both quantized and floating point.
    // This occurs when doing qi16 convolution, as bias is passed as a
    // non-quantized int64
    if (allTypesFp || allTypesQuantizedOrInt) return failure();

    Location loc = op->getLoc();
    SmallVector<Value> newOperands;

class NchwConvolutionToNhwcPass
    : public impl::NchwConvolutionToNhwcPassBase<NchwConvolutionToNhwcPass> {
 private:
  void runOnOperation() override;
};

// Rewrites NCHW convolution to NHWC.
// * Src dimension numbers: [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1]
// * Dst dimension numbers: [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
class RewriteNchwConvolutionToNhwc

        'b4': tf.Variable(tf.zeros([num_classes])),
    }

  @tf.function
  def __call__(self, data):
    """The Model definition."""
    x = tf.reshape(data, [-1, 28, 28, 1])

    # 2D convolution, with 'SAME' padding (i.e. the output feature map has
    # the same size as the input).

    # NOTE: The data/x/input is always specified in floating point precision.
    # output shape: [-1, 28, 28, 32]

}

// CHECK-LABEL: func @transposeConv2D_1x1_f32
func.func @transposeConv2D_1x1_f32(%input: tensor<1x64x28x28xf32>, %filter: tensor<1x1x64x64xf32>) -> tensor<1x64x28x28xf32> {
  // 1x1 convolution can be computed as a GEMM in NHWC data format.
  // CHECK: "tf.Conv2D"(%[[INPUT_TRANSPOSE:[0-9]*]], %arg1)
  // CHECK-SAME: data_format = "NHWC"
  %0 = "tf.Conv2D"(%input, %filter)
       {
         data_format = "NCHW",

// Make the 1D value $0 broadcastable with the shape of $1.
def MakeOneDimValueBroadcastable : NativeCodeCall<
  "MakeOneDimValueBroadcastable($_builder, $_loc, $0, $1.getType().cast<ShapedType>())">;

// Match convolution op with "NHWC" data format or matmul op.
def SupportedAffineOpMatcher : NativeCodeCall<
  "MatchSupportedAffineOp($_self, $0, $1, $2)">;

// Checks if a value can be symetrically quantized.

  pm.addNestedPass<func::FuncOp>(createFuseConvolutionPass());
  pm.addNestedPass<func::FuncOp>(createOptimizePass());
  // Conditionally enable below pass because this causes unfused convolutions
  // described in b/293149194. This problem is not replicated in
  // StableHLO Quantizer.
  if (add_fold_broadcast_pass) {
    pm.addNestedPass<func::FuncOp>(createFoldBroadcastPass());
  }

} {

// CHECK-LABEL: func @transposeConv2D_3x3_f16
func.func @transposeConv2D_3x3_f16(%input: tensor<1x28x28x64xf16>, %filter: tensor<3x3x64x64xf16>) -> tensor<1x26x26x64xf16> {
  // cuDNN prefers NCHW data format for spatial convolutions in f16 before
  // compute capability 7.0 (NVIDIA Tensor Cores).

  // CHECK: "tf.Conv2D"(%[[INPUT_TRANSPOSE:[0-9]*]], %arg1)
  // CHECK-SAME: data_format = "NCHW"
  %0 = "tf.Conv2D"(%input, %filter)
       {