FailureOr<SmallVector<int32_t>> CastI64ArrayToI32(
    const ArrayRef<int64_t> int64_array) {
  SmallVector<int32_t> int32_array{};
  int32_array.reserve(int64_array.size());

  for (const int64_t i64 : int64_array) {
    FailureOr<int32_t> cast_i32 = CastI64ToI32(i64);
    if (failed(cast_i32)) return failure();

    int32_array.push_back(*cast_i32);
  }
  return int32_array;
}

StringRef GetEntryFunctionName(TF::XlaCallModuleOp op) {

  // CHECK: %[[RESULT:.+]] = mhlo.add %[[X_NORMED]], %[[OFFSET_BCAST]] : tensor<4x256xf32>
  %0 = "mhlo.batch_norm_inference"(%x, %scale, %offset, %mean, %variance)
      {epsilon = 1.001000e-05 : f32, feature_index = 1 : i64} :
      (tensor<4x256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>,
        tensor<256xf32>) -> tensor<4x256xf32>
  // CHECK-DAG: return %[[RESULT]]
  func.return %0 : tensor<4x256xf32>

 // expected-error@below {{op is in dialect chlo which is not an accepted dialect}}
  %0 = "chlo.broadcast_add"(%arg0, %arg1) {broadcast_dimensions = array<i64: 3>} : (tensor<1x32x10x32xi32>, tensor<32xi32>) -> tensor<1x32x10x32xi32>
  func.return %0 : tensor<1x32x10x32xi32>

      %1 = "tf.XlaCallModule"(%0, %cst_0, %cst_1) <{Sout = [#tf_type.shape<1x4x4x3>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_2, _original_entry_function = "composite_fn_2", _tfl_quant_trait...

  EXPECT_EQ(l.size(), 1);
}

TaggedValue AddIntegers(TaggedValue args_, TaggedValue kwargs_) {
  auto& args = args_.tuple();
  // auto& kwargs = kwargs_.dict();
  return TaggedValue(args[0].i64() + args[1].i64());
}

TEST(ObjectTest, TestCast) {
  Integer i(3);
  auto result = Cast<String>(i);
  ASSERT_TRUE(!result.ok());
}

TEST(ObjectTest, TestCall) {
  TaggedValue add_func(AddIntegers);

    // CHECK:     return %[[RESULT]]
    func.return %0 : tensor<10xi32>
  }

  // CHECK-LABEL: func @foo

# CHECK-SAME: ({{%.*}}: tensor<*xf32>, {{%.*}}: tensor<*xi32> {tf._arg1_attr0 = "_arg1_attr0_value", tf._arg1_attr1 = 8.000000e+00 : f32}, {{%.*}}: tensor<*xi1>)
# CHECK-SAME: -> (tensor<*xf32> {tf._ret0_attr0 = 8 : i64, tf._ret0_attr1 = false}, tensor<*xi32>, tensor<*xi1> {tf._ret2_attr0 = !tf_type.variant, tf._ret2_attr1 = #tf_type.shape<128x1024>})

    // Pad the original state as necessary to fill the output shape.
    int pad = tensorflow::RNG_MAX_COUNTER_SIZE - counter_size;
    Type i64 = rewriter.getI64Type();
    RankedTensorType paddings_ty = RankedTensorType::get({1, 2}, i64);
    std::vector<int64_t> paddings_values = {0, pad};
    Value paddings = rewriter.create<ConstOp>(
        loc, DenseIntElementsAttr::get(paddings_ty, paddings_values));

    // CHECK-SAME:   _stablehlo_module_attrs = {}

    %0 = "tf.XlaCallModule"(%arg0, %arg1) {Sout = [#tf_type.shape<?>], dim_args_spec = [], _entry_function = @main_0, module = "", platforms = [], version = 5 : i64} : (tensor<10xi32>, tensor<10xi32>) -> tensor<10xi32>
    // CHECK: return %[[RESULT]]
    func.return %0 : tensor<10xi32>
  }

  // CHECK-LABEL: func private @_tf_func

      Sout = [#tf_type.shape<10>], dim_args_spec = [],
      _entry_function = @main_0,
      _stablehlo_module_attrs = { mhlo.num_partitions = 1 }, module = "",
      platforms = [], version = 5 : i64
    } : (tensor<10xf32>, tensor<10xf32>, tensor<6xi32>) -> (tensor<10xf32>, tensor<6xi32>)
    // CHECK: return %[[RESULT_CAST]], %[[RESULT]]#1 : tensor<10xf32>, tensor<6xi32>
    func.return %0#0, %0#1 : tensor<10xf32>, tensor<6xi32>