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()),

        input_output_alias.push_back(
            {new_argument_index, return_operands.size()});
      } else if (is_var_handle) {
        output_only_resources.push_back(
            {return_operands.size(),
             var_handle_shared_names[index - var_handles_start_idx]});
      }
      return_operands.push_back(resource.live_value);
      result_types.push_back(resource.live_value.getType());
    }

                         stride_i, pad_i_low, pad_i_high);
    temp_padding_values.push_back(scalar_to_rank1(pad_i_low));
    temp_padding_values.push_back(scalar_to_rank1(pad_i_high));
  }
  temp_padding_values.push_back(zero_rank1);
  temp_padding_values.push_back(zero_rank1);

  padding = CreateConstValue<int32_t>(
      builder, loc, /*shape=*/{num_dims - 2, 2},

  for (int i = 0; i < t.NumElements(); i++) {
    int64_t element = t.dtype() == DT_INT32
                          ? static_cast<int64_t>(t.flat<int32>()(i))
                          : t.flat<int64_t>()(i);
    result.push_back(element);
  }
  return result;
}

// Packages up the inputs to a Slice operation into an instance of
// `SliceInputs`.
StatusOrOptional<SliceInputs> GetSliceInputs(Node* slice) {
  const int kSliceInputIndex = 0;

    auto info_it = var_access_info.per_resource_info.find(resource);
    if (info_it == var_access_info.per_resource_info.end()) continue;
    device_var_reads_indices.push_back(info_it->second.execute_input_index);
    device_var_updates_indices.push_back(info_it->second.execute_output_index);
  }

  // Check that all resources are either read or written to.
  for (auto it : llvm::enumerate(var_access_info.new_operand_values)) {

        if (defining_op != nullptr &&
            defining_op->hasTrait<OpTrait::ConstantLike>()) {
          continue;
        }

        quantizable_values.push_back(input.get());
        aggregator_ids.push_back(
            (llvm::Twine(composite_function_name.value()) + "_arg_" +
             llvm::Twine(input.getOperandNumber()) + "_calibration_method_" +

      if (ops_with_tokens.insert(parent).second) {
        worklist.push_back(parent);
      }
      continue;
    }

    // For functions, get all the users and add them to the worklist.
    for (auto& user : symbol_map.getUsers(parent)) {
      if (ops_with_tokens.insert(user).second) {
        worklist.push_back(user);
      }
    }
  }

  // for the signature, for us, since we know the inputs, we can infer the time
  // steps.

  // Last output.
  outputs.push_back(last_output);
  output_types.push_back(last_output.getType());

  // Full sequences.
  outputs.push_back(final_output_full_sequences);
  output_types.push_back(final_output_full_sequences.getType());

  // All the rest: states, device.
  for (int i = 2; i < 5; ++i) {
    auto result_type =