contraction,
          "cannot fuse with the subsequent op as it's not supported by the "
          "target device.");
    }

    SmallVector<Location, 3> locations{contraction.getLoc(), bias_add.getLoc()};
    SmallVector<Attribute, 2> fused_ops{StringAttr::get(
        context, bias_add.getOperation()->getName().stripDialect())};

    // BiasAdd may or may not feed into an activation function.

    IntegerAttr group_key = rewriter.getI32IntegerAttr(0);

    auto const_group_size = rewriter.create<TF::ConstOp>(
        assign_group->getLoc(), assign_group.getResult(0).getType(),
        group_size);
    auto const_group_key = rewriter.create<TF::ConstOp>(
        assign_group->getLoc(), assign_group.getResult(1).getType(), group_key);
    rewriter.replaceAllUsesWith(assign_group.getResult(0), const_group_size);

  }

  auto lhs_new_concat = rewriter.create<mhlo::ConcatenateOp>(
      concat->getLoc(), concat.getType().clone(lhs_new_concat_shape),
      all_dot_lhs, rewriter.getI64IntegerAttr(lhs_batch_dim));
  auto rhs_new_concat = rewriter.create<mhlo::ConcatenateOp>(
      concat->getLoc(), concat.getType().clone(rhs_new_concat_shape),
      all_dot_rhs, rewriter.getI64IntegerAttr(rhs_batch_dim));

    // current op, so the result op has input with tensor type.
    if (auto cast_op = dyn_cast_or_null<CastOp>(input_op)) {
      Value input_to_cast = CastToNonDerivedType(rewriter, call_op.getLoc(),
                                                 cast_op, input_tfr_type);
      inputs->push_back(input_to_cast);
      if (failed(AddDerivedAttrs(rewriter, input_tfr_type,

        return false;
      }
    }
    rewriter.setInsertionPointAfter(op);
    auto q = rewriter.create<quantfork::QuantizeCastOp>(op->getLoc(), cast_type,
                                                        op.getResult());
    auto dq = rewriter.create<quantfork::DequantizeCastOp>(op->getLoc(),
                                                           expressed_type, q);
    quantize_op->setOperand(quantize_operand_num, dq.getResult());

    } else {
      return rewriter.notifyMatchFailure(bcast_op, "Unsupported element type.");
    }
    Value new_const_op =
        rewriter.create<mhlo::ConstantOp>(bcast_op.getLoc(), result);
    rewriter.replaceOp(bcast_op, {new_const_op});
    return success();
  }
};

using FoldBroadcastInDimBeforeMulOp =
    FoldBroadcastInDimBeforeBinaryElementwiseOp<mhlo::MulOp>;

  auto type = result.getType();
  result.setType(converter.convertType(result.getType()));
  auto new_value = converter.materializeArgumentConversion(
      rewriter, result.getLoc(), type, {result});
  rewriter.replaceAllUsesExcept(result, new_value, new_value.getDefiningOp());
}

// Wrap operands in an an unrealized cast to create a cast to buffer any type

        quant_type.castFromExpressedType(expressed_type);

    rewriter.setInsertionPointAfter(op);
    auto q = rewriter.create<quantfork::QuantizeCastOp>(
        op->getLoc(), quantized_type, op->getResult(0));
    auto dq = rewriter.create<quantfork::DequantizeCastOp>(op->getLoc(),
                                                           expressed_type, q);
    quantizable_op->setOperand(1, dq.getResult());
  }

 private:

      }
    }
  };

  for (int i = 0, e = func.getNumArguments(); i != e; ++i) {
    BlockArgument arg = func.getArgument(i);
    auto* arg_block = arg.getOwner();
    add_quantize_op(arg.getLoc(), arg.getType(), arg_block,
                    std::next(arg_block->begin(), i), arg, i);
  }

  return false;
}

bool PrepareQuantizePass::RemoveRedundantStats(func::FuncOp func) {

    rewriter.setInsertionPointAfter(weight_op);
    auto new_shape_const = rewriter.create<arith::ConstantOp>(
        weight_op->getLoc(), shape_spec_type, new_shape_const_attr);
    auto reshape_op = rewriter.create<TF::ReshapeOp>(
        weight_op->getLoc(), new_shape, weight_op->getResult(0),
        new_shape_const);
    op->setOperand(weight_operand_idx, reshape_op);