return success();
  }

  OpSet target_opset_;
};

// To calculate per-channel scale and offset, weight of depthwise was reshaped
// to [H, W, 1, InxMul]. After scale and offset has been calculated, this
// pattern gets called and restores the weight of depthwise back
// into [H, W, In, Mul]
class RestoreWeightShapePattern
    : public OpRewritePattern<TF::PartitionedCallOp> {

      strides: Sequence[int] = (1, 2, 2, 1),
      dilations: Sequence[int] = (1, 1, 1, 1),
      padding: str = 'SAME',
  ):
    class DepthwiseConvModel(module.Module):
      """A simple model with a single depthwise conv2d, bias and relu."""

      def __init__(self):
        self.out_channel_size = filter_shape[2] * filter_shape[3]

        # This ensures filters will have different value range per out channel