// Allocate output
    TF_Status* s = TF_NewStatus();
    int64_t dim = 1;
    size_t tensor_size_bytes = TF_DataTypeSize(TF_FLOAT);
    TF_Tensor* output = TF_AllocateOutput(
        /*context=*/ctx, /*index=*/0, /*dtype=*/TF_FLOAT, /*dims=*/&dim,
        /*num_dims=*/1, /*len=*/tensor_size_bytes, s);
    validate_tensor(output, &dim, 1, TF_FLOAT);

    // Set output to 3
    float values[1] = {3.0f};

  for (const auto &item : llvm::enumerate(dims)) {
    int64_t index = item.index();
    int64_t dim = item.value().getSExtValue();
    if (dim < 0 || dim > output_rank) {
      return emitOptionalError(location, "out of range broadcast dim");
    }
    if (is_broadcasted[dim]) {
      return emitOptionalError(location, "broadcast_dims has duplicates");
    }
    broadcast_shape[dim] = min_rank_ty.getDimSize(index);

//   ```
//
//   This operation requires that:
//
//   `-1-input.dims() <= dim <= input.dims()`
//
//   This operation is related to `squeeze()`, which removes dimensions of
//   size 1.
Status ExpandDims(AbstractContext* ctx, AbstractTensorHandle* const input,
                  AbstractTensorHandle* const dim,
                  AbstractTensorHandle** output, const char* name,

//
// May return NULL (and invoke the deallocator) if the provided data buffer
// (data, len) is inconsistent with a tensor of the given TF_DataType
// and the shape specified by (dima, num_dims).
TF_CAPI_EXPORT extern TF_Tensor* TF_NewTensor(
    TF_DataType, const int64_t* dims, int num_dims, void* data, size_t len,
    void (*deallocator)(void* data, size_t len, void* arg),
    void* deallocator_arg);

  llvm::SmallVector<int64_t, 4> dims;
  dims.reserve(shape.dim().size());
  for (const auto& dim : shape.dim()) {
    dims.push_back(dim.size() == kTFDynamicSize ? ShapedType::kDynamic
                                                : dim.size());
  }
  return mlir::TF::ShapeAttr::get(context, llvm::ArrayRef(dims));
}

      } else {
        const auto num_dims = tensor_shape.dim_size();
        std::unique_ptr<int64_t[]> dims(new int64_t[num_dims]);
        for (int i = 0; i < num_dims; ++i) {
          dims[i] = tensor_shape.dim(i).size();
        }
        TFE_OpSetAttrShape(op, attr_name, dims.get(), num_dims, status);
      }
    } break;
    case tensorflow::AttrValue::kFunc: {

      %1 = stablehlo.get_dimension_size %0, dim = 0 : (tensor<?x3xf32>) -> tensor<i32>
      %2 = stablehlo.reshape %1 : (tensor<i32>) -> tensor<1xi32>
      %3 = stablehlo.concatenate %2, %cst_0, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32>
      %4 = stablehlo.dynamic_broadcast_in_dim %arg2, %3, dims = [1] : (tensor<3xf32>, tensor<2xi32>) -> tensor<?x3xf32>

// $1.
// Requires $0 to be of RankedTensorType with rank greater than `dim` and the
// dimension should be known.
class GetDimSizeOfType<int dim> : NativeCodeCall<
  "GetScalarOfType(getElementTypeOrSelf($1), "
  "$0.getType().cast<RankedTensorType>().getDimSize(" # dim # "))">;

// Same as the above with i32 element type.
class GetDimSizeAsI32<int dim> : NativeCodeCall<
  "GetScalarOfType($_builder.getIntegerType(32), "

  return shaped_type.hasStaticShape();
}

bool HasStaticShapeAtDims(Value value, const ArrayRef<int> dims) {
  auto shaped_type = mlir::dyn_cast<ShapedType>(value.getType());
  if (!shaped_type || !shaped_type.hasRank()) return false;

  for (auto dim : dims) {
    if (shaped_type.isDynamicDim(dim)) return false;
  }
  return true;
}

Type CloneTypeWithNewElementType(Type old_type, Type element_type) {

    // Collect non-unit dims and corresponding dim in the input shape.
    SmallVector<int64_t, 4> input_carryover_dims;
    SmallVector<int64_t, 4> non_unit_dims;

    for (int i = 0; i < input_shape_extended.size(); i++) {
      int64_t dim = broadcast_shape[i];
      if (dim != 1) {
        non_unit_dims.push_back(dim);
        input_carryover_dims.push_back(input_shape_extended[i]);
      }