if (rhs_batch_idx.contains(i)) {
      if (rhs_only_batch) {
        rhs_eq.push_back(c);
        lhs_batch_dims.push_back(c);
        c++;
      } else {
        rhs_eq.push_back(lhs_batch_dims[batch_trace_idx]);
        batch_trace_idx++;
      }
    } else if (rhs_contract_idx.contains(i)) {
      absl::StrAppend(&rhs_eq,
                      std::string(1, lhs_contract_dims[contract_trace_idx]));

          rmax = std::max(std::abs(rmin), std::abs(rmax));
          rmin = -rmax;
        }
        TensorRangeSanityCheck(op, rmin, rmax);
        mins.push_back(rmin);
        maxs.push_back(rmax);
      }
      quant_type = quantfork::fakeQuantAttrsToType(
          op.getLoc(), num_bits, *op.getAxis(), mins, maxs, narrow_range,
          expressed, is_signed);

              op->getLoc(), unconstrainted_tensor_type, op->getOperand(i));
          variadic_operands.push_back(casted);
        }
        auto build_list_op = builder.create<BuildListOp>(
            op->getLoc(), list_type, variadic_operands);
        new_operands.push_back(build_list_op.getOut());
      } else {
        auto attr_name = compose_func.getArgAttrOfType<StringAttr>(

    result_types.push_back(op->getOpResult(2).getType());

    // Populate attributes.
    SmallVector<NamedAttribute, 4> attributes;
    // Activation will always be tanh.
    attributes.push_back(rewriter.getNamedAttr("fused_activation_function",
                                               rewriter.getStringAttr("TANH")));
    // cell_clip.
    attributes.push_back(

  std::vector<std::string> disabled_checks;
  for (auto attr : op.getDisabledChecks().getAsRange<StringAttr>()) {
    disabled_checks.push_back(attr.getValue().str());
  }
  std::vector<std::string> platforms;
  for (auto attr : op.getPlatforms().getAsRange<StringAttr>()) {
    platforms.push_back(attr.getValue().str());
  }
  TF_ASSIGN_OR_RETURN(
      auto loader,
      tensorflow::XlaCallModuleLoader::Create(

  } else {
    SmallVector<double> scales_vec;
    SmallVector<int64_t> zero_points_vec;
    for (auto elem : scales.getValues<float>()) scales_vec.push_back(elem);
    for (auto elem : zero_points.getValues<int32_t>())
      zero_points_vec.push_back(elem);
    elem_ty = quant::UniformQuantizedPerAxisType::get(
        flags, storage_type, expressed_type, scales_vec, zero_points_vec,

                      arg_locs);
  SmallVector<NamedAttribute> func_attrs;
  func_attrs.push_back(
      {StringAttr::get(context, "inputs"),
       StringAttr::get(context, absl::StrJoin(input_names, ","))});
  func_attrs.push_back(
      {StringAttr::get(context, "outputs"),
       StringAttr::get(context, absl::StrJoin(output_names, ","))});

    return split.emitOpError("different split lengths are not supported");
  }
  llvm::SmallVector<int64_t, 8> elem_shape;
  elem_shape.push_back(length);
  for (int64_t dim : buffer_type.getShape().drop_front()) {
    elem_shape.push_back(dim);
  }
  *elem_type = RankedTensorType::get(elem_shape, buffer_type.getElementType());
  return success();
}

// Tries to infer the tensor array element shape.

}

TEST_F(ArrayGradTest, ConcatV2Grad) {
  TensorShape shape({3, 2, 5});
  std::vector<Output> xs;
  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  xs.push_back(Placeholder(scope_, DT_FLOAT, Placeholder::Shape(shape)));
  auto axis = Const(scope_, 0);
  auto y = Concat(scope_, xs, axis);

    components.reserve(underlying_devices_.size());
    for (int j = 0; j < underlying_devices_.size(); ++j) {
      components.push_back(std::move(per_device_output_tensors[j][i]));
    }
    if (expected_output_shapes[i].IsFullyDefined()) {
      per_device_outputs.push_back(ParallelTensor::FromTensorHandles(
          *this, std::move(components),
          absl::Span<const int64_t>(expected_output_shapes[i].dim_sizes()),