broadcast_multiplier);
    auto broadcast_rhs = broadcastToFeatureDim(
        bn_op.getLoc(), input_type, rhs, shape_value, feature_dim, rewriter);
    rewriter.replaceOpWithNewOp<mhlo::AddOp>(bn_op, lhs, broadcast_rhs);

    return success();
  }
};

class UnfuseMhloBatchNormPass
    : public PassWrapper<UnfuseMhloBatchNormPass, OperationPass<func::FuncOp>> {
 public:

          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),
          rewriter.getStringAttr("TPU"));
    }

    constexpr int64_t kDefaultCostEstimate = 1000000;
    rewriter.replaceOpWithNewOp<TF::XlaHostComputeOp>(
        op, op.getResultTypes(), op.getInputs(),
        /*ancestors=*/rewriter.getArrayAttr({}),
        rewriter.getArrayAttr(shape_attrs),
        /*shape_inference_graph=*/

  const auto matmul_type = RankedTensorType::get({rows, cols}, element_type);

  if (lhs_dims == 2 && rhs_dims == 2) {
    // When both inputs are matrices, just replace the op with a matmul op.
    rewriter.replaceOpWithNewOp<TF::MatMulOp>(op, matmul_type,
                                              /*a=*/input_lhs,
                                              /*b=*/input_rhs,

        quantized_attr = quant::QuantizeLegacy(attr, qtype.getValue());
      } else {
        quantized_attr = quant::Quantize(attr, qtype.getValue());
      }
      if (quantized_attr) {
        rewriter.replaceOpWithNewOp<QConstOp>(op, qtype, quantized_attr);
        return success();
      }
    }
    return failure();
  }

 private:
  bool legacy_float_scale_;
};

    // Update the scale in the quantize op.
    auto new_qtype = RescaleQtype(q_op.getQtype(), gamma_cst);
    if (!new_qtype) return failure();
    rewriter.replaceOpWithNewOp<TFL::QuantizeOp>(q_op, new_qtype.getValue(),
                                                 new_filter, new_qtype);

    // If bias isn't None, it needs to be multiplied as well.

        mlir::dyn_cast<TensorType>(ToLegalType(op.getOutput().getType())));
    if (failed(dense_attr_or)) {
      op->emitError("failed to get DenseElementAttr.");
      return failure();
    }

    rewriter.replaceOpWithNewOp<TF::ConstOp>(
        op, ToLegalType(op.getOutput().getType()), *dense_attr_or);
    return success();
  }
};

struct ConvertTFQuantTypes
    : public impl::ConvertTFQuantTypesBase<ConvertTFQuantTypes> {

                                        ~strided_slice_op.getNewAxisMask();

    auto attribute_type = rewriter.getIntegerType(64);
    rewriter.replaceOpWithNewOp<TF::StridedSliceOp>(
        op, strided_slice_op.getType(), reshape, strided_slice_op.getBegin(),
        strided_slice_op.getEnd(), strided_slice_op.getStrides(),

      rewriter.replaceOpWithNewOp<arith::ConstantOp>(const_op,
                                                     const_op.getValue());
      return success();
    } else if (TFL::NoValueOp::isBuildableWith(const_op.getValue(),
                                               const_op.getType())) {
      rewriter.replaceOpWithNewOp<NoValueOp>(
          const_op, rewriter.getNoneType(),

  auto new_mean_op =
      rewriter.create<TFL::MeanOp>(mean_op->getLoc(), new_output_type, input,
                                   mean_op.getAxis(), mean_op.getKeepDims());

  // Insert a requant op.
  rewriter.replaceOpWithNewOp<TFL::QuantizeOp>(
      mean_op, output_type, new_mean_op, mlir::TypeAttr::get(output_type));
  return success();
}

}  // namespace tac
}  // namespace TFL

  // canonical.
  if (cluster_vals.size() == num_results) return failure();

  // Rewrite return op in the cluster.
  rewriter.setInsertionPoint(return_op);
  auto new_return =
      rewriter.replaceOpWithNewOp<tf_device::ReturnOp>(return_op, cluster_vals);

  // Rewrite the cluster op.
  rewriter.setInsertionPoint(op);
  auto new_op = rewriter.create<tf_device::ClusterOp>(