message: Rollback change to FakeQuantWithMinMaxVars. The change caused a small increase in error in //research/soapbox/app/golden7/test:tpu_model_validation_test. We'd like to understand the reason for the change.
diff: @@ -13,7 +13,6 @@ See the License for the specific language governing permissions and
  limitations under the License.
  ==============================================================================*/
  -#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
  #include "tensorflow/compiler/tf2xla/xla_helpers.h"
  #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
  #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
  @@ -223,7 +222,53 @@ class FakeQuantWithMinMaxArgsGradOp : public XlaOpKernel {
  REGISTER_XLA_OP(Name("FakeQuantWithMinMaxArgsGradient"),
  FakeQuantWithMinMaxArgsGradOp);
  -REGISTER_XLA_OP(Name("FakeQuantWithMinMaxVars"), MlirXlaOpKernel);
  +class FakeQuantWithMinMaxVarsOp : public XlaOpKernel {
  + public:
  + explicit FakeQuantWithMinMaxVarsOp(OpKernelConstruction* ctx)
  + : XlaOpKernel(ctx) {
  + OP_REQUIRES_OK(ctx, ctx->GetAttr("num_bits", &num_bits_));
  + OP_REQUIRES(ctx, num_bits_ >= 2 && num_bits_ <= 16,
  + errors::InvalidArgument("num_bits is out of range, expected "
  + "between 2 and 16, was: ",
  + num_bits_));
  + OP_REQUIRES_OK(ctx, ctx->GetAttr("narrow_range", &narrow_range_));
  + quant_min_ = narrow_range_ ? 1 : 0;
  + quant_max_ = (1 << num_bits_) - 1;
  + }
  +
  + void Compile(XlaOpKernelContext* ctx) override {
  + xla::XlaBuilder* b = ctx->builder();
  + xla::XlaOp input = ctx->Input(0);
  + const DataType data_type = ctx->input_type(0);
  + xla::XlaOp input_min = ctx->Input(1);
  + xla::XlaOp input_max = ctx->Input(2);
  +
  + if (ctx->compiler()->options().allow_cpu_custom_calls &&
  + ctx->compiler()->options().custom_fake_quant_op_calls) {
  + xla::XlaOp custom_call_output =
  + b->ReportErrorOrReturn(BuildFakeQuantCustomCall(
  + b, input, input_min, input_max, num_bits_, narrow_range_));
  + ctx->SetOutput(0, custom_call_output);
  + return;
  + }
  +
  + xla::XlaOp nudged_input_min, nudged_input_max, input_scale;
  + XlaNudge(b, data_type, input_min, input_max, quant_min_, quant_max_,
  + &nudged_input_min, &nudged_input_max, &input_scale);
  +
  + xla::XlaOp output = Quantize(b, input, data_type, nudged_input_min,
  + nudged_input_max, input_scale);
  + ctx->SetOutput(0, output);
  + }
  +
  + private:
  + int num_bits_;
  + bool narrow_range_;
  + float quant_min_;
  + float quant_max_;
  +};
  +
  +REGISTER_XLA_OP(Name("FakeQuantWithMinMaxVars"), FakeQuantWithMinMaxVarsOp);
  class FakeQuantWithMinMaxVarsGradOp : public XlaOpKernel {
  public:
  