ToBfloat16Type(
          RankedTensorType::get({2, 2}, Float8E4M3FNType::get(context.get()))),
      RankedTensorType::get({2, 2}, Float8E4M3FNType::get(context.get())));
  EXPECT_EQ(ToBfloat16Type(
                RankedTensorType::get({2, 2}, Float16Type::get(context.get()))),
            RankedTensorType::get({2, 2}, Float16Type::get(context.get())));
  EXPECT_EQ(ToBfloat16Type(RankedTensorType::get(

  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());
  auto weight_type = RankedTensorType::get(weight_shape, builder->getF32Type());
  auto bias_type = RankedTensorType::get(bias_shape, builder->getF32Type());
  auto projection_type =
      RankedTensorType::get(projection_shape, builder->getF32Type());
  auto layer_norm_scale_type =
      RankedTensorType::get(layer_norm_scale, builder->getF32Type());

  auto output = op->getResult(0);
  auto output_type = mlir::dyn_cast_or_null<RankedTensorType>(output.getType());
  if (!output_type) return failure();

  // bias should be a vector sized of the last output dim.
  int64_t num_units = output_type.getDimSize(output_type.getRank() - 1);
  auto bias_type =
      mlir::RankedTensorType::get({num_units}, output_type.getElementType());

  mlir::DenseElementsAttr bias_attr;

DenseIntElementsAttr GetI64ElementsAttr(ArrayAttr attr) {
  RankedTensorType ty =
      RankedTensorType::get(static_cast<int64_t>(attr.size()),
                            IntegerType::get(attr.getContext(), 64));
  return DenseIntElementsAttr::get(ty, attr.getValue());
}

DenseIntElementsAttr GetI64ElementsAttr(ArrayRef<int64_t> values,
                                        Builder* builder) {
  RankedTensorType ty = RankedTensorType::get(

    RankedTensorType u64_scalar = RankedTensorType::get({}, u64);
    Value step_size = rewriter.create<ConstOp>(loc, GetScalarOfType(u64, 256));
    Value increment =
        rewriter.create<MulOp>(loc, u64_scalar, step_size, rng_op.getDelta());

    // Increment the counter.
    SmallVector<Value, 4> pack_args;
    RankedTensorType word_u64_type = RankedTensorType::get({}, u64);

Value getShapeValue(Location loc, Value operand, PatternRewriter &rewriter) {
  RankedTensorType resultType =
      mlir::dyn_cast<RankedTensorType>(operand.getType());
  return rewriter.create<shape::ShapeOfOp>(
      loc,
      RankedTensorType::get(/*shape=*/{resultType.getRank()},
                            rewriter.getIndexType()),
      operand);
}

}

static RankedTensorType GetPermutedTensorTypeHelper(RankedTensorType type,
                                                    ArrayRef<int64_t> perm,
                                                    bool isInvert) {
  SmallVector<int64_t, 4> permutedShape = applyPermutation(
      type.getShape(), isInvert ? invertPermutationVector(perm) : perm);
  return RankedTensorType::get(permutedShape, type.getElementType());
}

namespace TF {

class IdentityOp;
class IdentityNOp;

// Returns the RankedTensorType for the given operand. TensorFlow constant ops
// may have non-static shape because the shape is not propagated during constant
// folding. If the defining op for the given operand is a constant op, this
// routine uses the constant op's attribute to get the actual shape.
RankedTensorType GetRankedTensorTypeForOperand(Value operand) {
  DenseElementsAttr attr;

  // Setup correct weights.
  RankedTensorType weight_type =
      mlir::cast<RankedTensorType>(weight_kernel.getType());
  if (weight_type.getRank() != 2)
    return func_op.emitError() << "The weight should be rank of 2";

  Value transposed_weight_kernel =
      Transpose2D(builder, weight_kernel, weight_type, func_op.getLoc());

  RankedTensorType recurrent_kernel_type =

    return false;
  }

  auto opr1_type =
      llvm::dyn_cast_or_null<RankedTensorType>(op->getOperand(0).getType());
  auto opr2_type =
      llvm::dyn_cast_or_null<RankedTensorType>(op->getOperand(1).getType());
  auto res_type =
      llvm::dyn_cast_or_null<RankedTensorType>(op->getResult(0).getType());

  if (!(opr1_type && opr2_type && res_type)) {
    return false;
  }