// element is arbitrarily considered to be 4 bytes.
constexpr int64_t kAssumedNumBytesPerElem = 4;

int64_t GetSizeOfIntOrFloatConst(TF::ConstOp const_op) {
  const Type dtype = const_op.getDtype();
  const ElementsAttr const_value = const_op.getValue();

  const auto bytes_per_elem =
      static_cast<int64_t>(dtype.getIntOrFloatBitWidth() / CHAR_BIT);

  return bytes_per_elem * const_value.getNumElements();

def : Pat<(TF_RandomUniformOp:$old $shape, $seed, $seed2),
          (MHLO_RngOp
            (MHLO_ConstantOp
              (NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 0.0)">)),
            (MHLO_ConstantOp
              (NativeCodeCall<"$_builder.getFloatAttr(old.getDtype(), 1.0)">)),
            (CastValueToI64 $old, $shape),
            MHLO_RngDistributionValue<"UNIFORM">),
          [(IsShapedTensor $shape)]>;

class TapeTensor {
 public:
  explicit TapeTensor(AbstractTensorHandle* handle);
  TapeTensor(const TapeTensor& other);
  ~TapeTensor();

  int64_t GetID() const;
  tensorflow::DataType GetDType() const;

  AbstractTensorHandle* ZerosLike() const;

  AbstractTensorHandle* GetHandle() const;

 private:
  AbstractTensorHandle* handle_;
};

// A tracing/immediate-execution agnostic tape.

  auto cast_op =
      OpBuilder(old_val.getDefiningOp())
          .create<tensor::CastOp>(old_val.getLoc(), old_val.getType(), new_val);
  old_val.replaceAllUsesWith(cast_op);
}

void ReplaceAllUsesExceptTerminator(Value old_val, Value new_val) {
  if (old_val.getType() == new_val.getType()) {
    old_val.replaceAllUsesWith(new_val);
    return;
  }
  Operation* old_op = old_val.getDefiningOp();

  sources.reserve(output_tensors.size());
  for (const TapeTensor& output_tensor : output_tensors) {
    // Ownership of `aid` transferred to CallBackwardFunction below.
    Gradient* aid;
    if (output_tensor.GetDType() == tensorflow::DT_VARIANT) {
      // Note: Needs to be zeros rather than ones since there's currently no
      // ones_like for variants.
      aid = output_tensor.ZerosLike();
    } else {

  handle_ = other.handle_;
  handle_->Ref();
}
TapeTensor::~TapeTensor() { handle_->Unref(); }

int64_t TapeTensor::GetID() const { return ToId(handle_); }

tensorflow::DataType TapeTensor::GetDType() const {
  return handle_->DataType();
}
AbstractTensorHandle* TapeTensor::GetHandle() const { return handle_; }

AbstractTensorHandle* TapeTensor::ZerosLike() const { return nullptr; }

class TapeVSpace

                CallableInfo callable = request.getInterceptedCallable();
                CallableKindInfo kind = callable.getKind();
                if (kind == CallableKindInfo.AFTER_CONSTRUCTOR) {
                    Type constructedType = request.getInterceptedCallable().getOwner().getType();
                    String typeKey = implementationName + ":" + constructedType;

    assertTrue(
        from.getType() + " is expected to be assignable to " + to.getType(),
        to.isSupertypeOf(from));
    assertTrue(
        to.getType() + " is expected to be a supertype of " + from.getType(),
        to.isSupertypeOf(from));
    assertTrue(
        from.getType() + " is expected to be a subtype of " + to.getType(), from.isSubtypeOf(to));
  }

    private static String getFqName(CallInterceptionRequest request) {
        String propertyName = request.getRequestExtras().getByType(PropertyUpgradeRequestExtra.class).get().getPropertyName();
        String containingType = request.getInterceptedCallable().getOwner().getType().getClassName();
        return containingType + "#" + propertyName;
    }

  auto result_type =
      OpTrait::util::getBroadcastedType(x.getType(), y.getType());
  if (!result_type) {
    if (incompatible_shape_error.getValue()) {
      mlir::emitError(loc, "non-broadcastable operands");
    } else {
      return UnrankedTensorType::get(builder->getI1Type());
    }
  }

  auto ranked_type = mlir::dyn_cast<RankedTensorType>(result_type);