}
  float& arg0(size_t dim0, size_t dim1) {
    return (*static_cast<float(*)[1][2]>(
        arg_data(0)))[dim0][dim1];
  }
  const float* arg0_data() const {
    return static_cast<const float*>(arg_data(0));
  }
  const float& arg0(size_t dim0, size_t dim1) const {
    return (*static_cast<const float(*)[1][2]>(
        arg_data(0)))[dim0][dim1];
  }
  int arg0_size() const {

    return a;
  }
  int64_t rank = a.getRank();
  SmallVector<int64_t, 4> dims;
  dims.resize(rank);
  for (int i = 0, e = rank; i != e; i++) {
    int64_t dim0 = a.getDimSize(i);
    int64_t dim1 = b.getDimSize(i);
    dims[i] = (dim0 == ShapedType::kDynamic) ? dim1 : dim0;
  }
  return RankedTensorType::get(dims, a.getElementType());
}
}  // namespace detail

    absl::string_view node_shapes_str) {
  std::vector<int> dims;
  for (absl::string_view dim_str : absl::StrSplit(node_shapes_str, ',')) {
    // Treats empty input shape as scalar
    if (dim_str.empty()) continue;
    if (dim_str == "?") {
      dims.push_back(-1);
      continue;
    }
    int size;
    TF_RET_CHECK(absl::SimpleAtoi(dim_str, &size));
    dims.push_back(size);
  }
  return dims;
}

    auto dim_nums = first_dot.getDotDimensionNumbers();
    return llvm::is_contained(dim_nums.getLhsBatchingDimensions(), dim) ||
           llvm::is_contained(dim_nums.getLhsContractingDimensions(), dim);
  };

  // dot_general outputs are always in the
  //   [batch dims, LHS other dims, RHS other dims]
  // layout, so the new concat dim is where the n-th (base-0 counting) LHS other

  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);

//
// 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), "