%1 = tfr.get_shape %0 -> !shape.shape
  %2 = shape.to_extent_tensor %1 : !shape.shape -> tensor<2xindex>
  func.return %2 : tensor<2xindex>

// CHECK-NEXT: %[[e:.*]] = shape.const_shape [1, 2] : tensor<2xindex>
// CHECK-NEXT: return %[[e]] : tensor<2xindex>
}

// CHECK-LABEL: equal
func.func @equal() -> (i1, i1, i1, i1) {
  %0 = tfr.constant f32 -> !tfr.attr
  %1 = tfr.constant f32 -> !tfr.attr

  // TODO(jpienaar): Uncomment below when fallout from https://reviews.llvm.org/D83194 fixed.
  // NOT-CHECK-DAG:  %[[LHS_SHAPE:.+]] = shape.shape_of %arg0
  // NOT-CHECK-DAG:  %[[RHS_SHAPE:.+]] = shape.const_shape [1]
  // NOT-CHECK-NEXT: %[[WITNESS:.+]] = shape.cstr_broadcastable %[[LHS_SHAPE]], %[[RHS_SHAPE]]
  // NOT-CHECK-NEXT: shape.assuming %[[WITNESS]] -> (tensor<?xi1>) {

    llvm::SmallVector<int64_t, 4> const_shape;
    if (attr_shape.isSplat()) {
      const_shape.assign(
          attr_shape.getNumElements(),
          x_type.getDimSize((*attr_shape.begin()).getSExtValue()));
    } else {
      const_shape.reserve(attr_shape.getNumElements());
      for (const auto &dim : attr_shape)
        const_shape.push_back(x_type.getDimSize(dim.getSExtValue()));
    }

    }

    auto lhs_shape = mlir::cast<ShapedType>(lhs.getType()).getShape();
    auto rhs_shape = mlir::cast<ShapedType>(rhs.getType()).getShape();
    auto cond_shape = mlir::cast<ShapedType>(cond.getType()).getShape();

    if (lhs_shape == rhs_shape && cond_shape == lhs_shape) {
      return failure();
    }

    // Calculate the broadcasted shape.
    SmallVector<int64_t, 4> broadcasted_shape;

    int arg_num = compilation_result->input_mapping[i];
    CHECK_GE(arg_num, missing_ctx_input_prefix);
    const xla::Shape& device_shape = compilation_result->xla_input_shapes[i];
    const xla::Shape& host_shape =
        xla::ShapeUtil::DeviceShapeToHostShape(device_shape);

    auto resource_var_it = resource_vars.find(arg_num);
    bool is_resource_variable = resource_var_it != resource_vars.end();

  bool has_static_cond_shape = false;
  bool has_static_operand_shape = false;
  ArrayRef<int64_t> cond_shape;
  ArrayRef<int64_t> operand_shape;

  if (auto shaped_type = cond.getType().dyn_cast<ShapedType>()) {
    if (shaped_type.hasStaticShape()) {
      has_static_cond_shape = true;
      cond_shape = shaped_type.getShape();
    }
  }
  if (auto shaped_type = operand_type.dyn_cast<ShapedType>()) {

    // Now combine x and y into the index data structure needed for gather.
    Shape concat_shape({2, num_diags, max_diag_len});
    Value start_indices = rewriter.create<ConcatenateOp>(
        loc,
        tensorflow::GetTypeFromTFTensorShape(concat_shape,
                                             rewriter.getIntegerType(32)),
        mlir::ValueRange({n_plus_y, n_plus_x}),