Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 14 additions & 14 deletions tensorflow_text/core/kernels/constrained_sequence.cc
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ ScoreAccessor::ScoreAccessor(const Tensor &score_tensor,
data_ = score_tensor.flat<float>().data();
if (lengths_tensor.dtype() == DT_INT64) {
use_long_lengths_ = true;
long_lengths_ = lengths_tensor.flat<int64>().data();
long_lengths_ = lengths_tensor.flat<int64_t>().data();
} else {
use_long_lengths_ = false;
lengths_ = lengths_tensor.flat<int>().data();
Expand Down Expand Up @@ -66,7 +66,7 @@ float ScoreAccessor::GetScore(int batch_idx, int step_idx,
return data_[batch_offset_ * batch_idx + step_offset_ * step_idx + score_idx];
}

int64 ScoreAccessor::GetLength(int batch_idx) const {
int64_t ScoreAccessor::GetLength(int batch_idx) const {
DCHECK_LE(batch_idx, batch_size_);
if (use_long_lengths_) {
return long_lengths_[batch_idx];
Expand All @@ -82,18 +82,18 @@ bool ScoreAccessor::has_explicit_batch() const { return has_explicit_batch_; }

// Perform Viterbi analysis on a single batch item.
void ViterbiAnalysis(
const ScoreAccessor &scores,
const tensorflow::TTypes<const float>::Matrix &transition_weights,
const tensorflow::TTypes<const bool>::Matrix &allowed_transitions,
const ScoreAccessor& scores,
const tensorflow::TTypes<const float>::Matrix& transition_weights,
const tensorflow::TTypes<const bool>::Matrix& allowed_transitions,
const int batch, bool use_log_space, bool use_start_end_states,
int32 *output_data) {
int32_t* output_data) {
VLOG(2) << "Analyzing batch " << batch;
const bool has_transition_weights = transition_weights.size() != 0;
const bool has_allowed_transitions = allowed_transitions.size() != 0;
const int num_states = scores.num_scores();
const int out_of_bounds_index = num_states;

int64 num_steps = scores.GetLength(batch);
int64_t num_steps = scores.GetLength(batch);

// Create two vectors to hold scores. These will be bound to referents later
// so the names here are somewhat irrelevant.
Expand Down Expand Up @@ -344,31 +344,31 @@ void ViterbiAnalysis(
if (best_source_state == kErrorState) {
// If the best source is an error state, the path is unknowable. Report
// error states for the whole sequence.
for (int64 i = 0; i < scores.GetLength(batch); ++i) {
for (int64_t i = 0; i < scores.GetLength(batch); ++i) {
output_data[i] = kErrorState;
}
} else {
// If the best source is a 'real' state, report the state path.
int steps_to_report = scores.GetLength(batch);
int previous_state = best_source_state;
for (int64 i = steps_to_report - 1; i >= 0; --i) {
for (int64_t i = steps_to_report - 1; i >= 0; --i) {
output_data[i] = previous_state;
previous_state = backpointers[i][previous_state];
}
}
}

void GreedyAnalysis(
const ScoreAccessor &scores,
const tensorflow::TTypes<const float>::Matrix &transition_weights,
const tensorflow::TTypes<const bool>::Matrix &allowed_transitions,
const ScoreAccessor& scores,
const tensorflow::TTypes<const float>::Matrix& transition_weights,
const tensorflow::TTypes<const bool>::Matrix& allowed_transitions,
int batch, bool use_log_space, bool use_start_end_states,
int32 *output_data) {
int32_t* output_data) {
const bool has_transition_weights = transition_weights.size() != 0;
const bool has_allowed_transitions = allowed_transitions.size() != 0;
const int num_states = scores.num_scores();
const int out_of_bounds_index = num_states;
int64 num_steps = scores.GetLength(batch);
int64_t num_steps = scores.GetLength(batch);

for (int step = 0; step < num_steps; ++step) {
// Do final step calculations if this is the final step in the sequence
Expand Down
20 changes: 10 additions & 10 deletions tensorflow_text/core/kernels/constrained_sequence.h
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ class ScoreAccessor {
// Get a score out of the data tensor.
float GetScore(int batch_idx, int step_idx, int score_idx) const;

int64 GetLength(int batch_idx) const;
int64_t GetLength(int batch_idx) const;

int batch_size() const;
int num_steps() const;
Expand All @@ -43,7 +43,7 @@ class ScoreAccessor {

// A pointer into the underlying data of the lengths tensor. Not owned.
const int *lengths_;
const int64 *long_lengths_;
const int64_t* long_lengths_;

// Whether the passed lengths tensor is int32 or int64.
bool use_long_lengths_;
Expand Down Expand Up @@ -72,19 +72,19 @@ class ScoreAccessor {

// Perform Viterbi analysis on a single batch item.
void ViterbiAnalysis(
const ScoreAccessor &scores,
const tensorflow::TTypes<const float>::Matrix &transition_weights,
const tensorflow::TTypes<const bool>::Matrix &allowed_transitions,
const ScoreAccessor& scores,
const tensorflow::TTypes<const float>::Matrix& transition_weights,
const tensorflow::TTypes<const bool>::Matrix& allowed_transitions,
const int batch, bool use_log_space, bool use_start_end_states,
int32 *output_data);
int32_t* output_data);

// Perform a greedy analysis on a single batch item.
void GreedyAnalysis(
const ScoreAccessor &scores,
const tensorflow::TTypes<const float>::Matrix &transition_weights,
const tensorflow::TTypes<const bool>::Matrix &allowed_transitions,
const ScoreAccessor& scores,
const tensorflow::TTypes<const float>::Matrix& transition_weights,
const tensorflow::TTypes<const bool>::Matrix& allowed_transitions,
int batch, bool use_log_space, bool use_start_end_states,
int32 *output_data);
int32_t* output_data);

} // namespace text
} // namespace tensorflow
Expand Down
31 changes: 15 additions & 16 deletions tensorflow_text/core/kernels/constrained_sequence_kernel.cc
Original file line number Diff line number Diff line change
Expand Up @@ -57,13 +57,12 @@ namespace {

// Validate that a given constraint tensor is the proper shape (dimension
// 2, with shape [num_states + 1, num_states + 1].
absl::Status ValidateConstraintTensor(const Tensor &tensor,
absl::Status ValidateConstraintTensor(const Tensor& tensor,
const int num_states,
const bool use_start_end_states,
const string &name) {
const std::string& name) {
if (tensor.shape().dims() != 2) {
return InvalidArgument(
tensorflow::strings::StrCat(name, " must be of rank 2"));
return InvalidArgument(absl::StrCat(name, " must be of rank 2"));
}
int expected_size = use_start_end_states ? num_states + 1 : num_states;
if (tensor.shape().dim_size(0) != expected_size) {
Expand Down Expand Up @@ -110,7 +109,7 @@ class ConstrainedSequence : public OpKernel {
const int num_scores = scores.num_scores();

OP_REQUIRES(context, lengths_tensor.NumElements() == batch_size,
InvalidArgument(tensorflow::strings::StrCat(
InvalidArgument(absl::StrCat(
"There should be exactly one length for every batch "
"element. Found ",
lengths_tensor.NumElements(),
Expand All @@ -124,7 +123,7 @@ class ConstrainedSequence : public OpKernel {
int max_length = 0;
int total_length = 0;
for (int i = 0; i < batch_size; ++i) {
int64 length = scores.GetLength(i);
int64_t length = scores.GetLength(i);
total_length += length;
if (length > max_length) {
max_length = length;
Expand Down Expand Up @@ -182,7 +181,7 @@ class ConstrainedSequence : public OpKernel {
Tensor *output;
OP_REQUIRES_OK(context, context->allocate_output(
0, TensorShape({total_length}), &output));
int32 *output_data = output->flat<int32>().data();
int32_t* output_data = output->flat<int32_t>().data();

Tensor *offsets;
OP_REQUIRES_OK(context, context->allocate_output(
Expand All @@ -192,7 +191,7 @@ class ConstrainedSequence : public OpKernel {

for (int batch = 0; batch < batch_size; ++batch) {
int step_offset = offset_data[batch];
int64 num_steps = scores.GetLength(batch);
int64_t num_steps = scores.GetLength(batch);
offset_data[batch + 1] = step_offset + num_steps;
if (use_viterbi_) {
DoViterbiAnalysis(transition_weights, allowed_transitions, batch,
Expand All @@ -207,18 +206,18 @@ class ConstrainedSequence : public OpKernel {
private:
// Perform Viterbi analysis on a single batch item.
void DoViterbiAnalysis(
const tensorflow::TTypes<const float>::Matrix &transition_weights,
const tensorflow::TTypes<const bool>::Matrix &allowed_transitions,
const int batch, const ScoreAccessor &scores, int32 *output_data) {
const tensorflow::TTypes<const float>::Matrix& transition_weights,
const tensorflow::TTypes<const bool>::Matrix& allowed_transitions,
const int batch, const ScoreAccessor& scores, int32_t* output_data) {
ViterbiAnalysis(scores, transition_weights, allowed_transitions, batch,
use_log_space_, use_start_end_states_, output_data);
}

// Perform a greedy analysis on a single batch item.
void DoGreedyAnalysis(
const tensorflow::TTypes<const float>::Matrix &transition_weights,
const tensorflow::TTypes<const bool>::Matrix &allowed_transitions,
int batch, const ScoreAccessor &scores, int32 *output_data) {
const tensorflow::TTypes<const float>::Matrix& transition_weights,
const tensorflow::TTypes<const bool>::Matrix& allowed_transitions,
int batch, const ScoreAccessor& scores, int32_t* output_data) {
GreedyAnalysis(scores, transition_weights, allowed_transitions, batch,
use_log_space_, use_start_end_states_, output_data);
}
Expand Down Expand Up @@ -251,8 +250,8 @@ class ConstrainedSequence : public OpKernel {
.TypeConstraint<int64>("Tsplits"), \
ConstrainedSequence<Tin, int64>)

REGISTER_KERNELS(int32);
REGISTER_KERNELS(int64);
REGISTER_KERNELS(int32_t);
REGISTER_KERNELS(int64_t);

#undef REGISTER_KERNELS

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,8 @@ TEST_F(ConstrainedSequenceInputValidationTest, WorksWithInt64InputLengths) {
}});

// Add the sequence_lengths input.
std::vector<int64> input_lengths({1, 1, 1});
AddInputFromArray<int64>(TensorShape({3}), input_lengths);
std::vector<int64_t> input_lengths({1, 1, 1});
AddInputFromArray<int64_t>(TensorShape({3}), input_lengths);

// Add the allowed_transitions input.
AddInputFromArray<bool>(TensorShape({5, 5}),
Expand All @@ -99,8 +99,8 @@ TEST_F(ConstrainedSequenceInputValidationTest, WorksWithInt64InputLengths) {
// The third sequence's highest score is 0, which is ok.

// Validate the output.
std::vector<int32> expected_transitions({1, 3, 0});
std::vector<int64> expected_offsets({0, 1, 2, 3});
std::vector<int32_t> expected_transitions({1, 3, 0});
std::vector<int64_t> expected_offsets({0, 1, 2, 3});

// Validate the output.
EXPECT_THAT(*GetOutput(0), VectorEq(expected_transitions));
Expand Down
15 changes: 8 additions & 7 deletions tensorflow_text/core/kernels/disjoint_set_forest_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -62,10 +62,10 @@ class DisjointSetForestTest : public ::testing::Test {
};

using Forests = ::testing::Types<
DisjointSetForest<uint8, false>, DisjointSetForest<uint8, true>,
DisjointSetForest<uint16, false>, DisjointSetForest<uint16, true>,
DisjointSetForest<uint32, false>, DisjointSetForest<uint32, true>,
DisjointSetForest<uint64, false>, DisjointSetForest<uint64, true>>;
DisjointSetForest<uint8_t, false>, DisjointSetForest<uint8_t, true>,
DisjointSetForest<uint16_t, false>, DisjointSetForest<uint16_t, true>,
DisjointSetForest<uint32_t, false>, DisjointSetForest<uint32_t, true>,
DisjointSetForest<uint64_t, false>, DisjointSetForest<uint64_t, true>>;
TYPED_TEST_SUITE(DisjointSetForestTest, Forests);

TYPED_TEST(DisjointSetForestTest, DefaultEmpty) {
Expand Down Expand Up @@ -111,12 +111,13 @@ TYPED_TEST(DisjointSetForestTest, Populated) {
// merged set can be controlled.
class DisjointSetForestNoUnionByRankTest : public ::testing::Test {
protected:
using Forest = DisjointSetForest<uint32, false>;
using Forest = DisjointSetForest<uint32_t, false>;

// Expects that the roots of the |forest| match |expected_roots|.
void ExpectRoots(const std::vector<uint32> &expected_roots, Forest *forest) {
void ExpectRoots(const std::vector<uint32_t>& expected_roots,
Forest* forest) {
ASSERT_EQ(expected_roots.size(), forest->size());
for (uint32 i = 0; i < forest->size(); ++i) {
for (uint32_t i = 0; i < forest->size(); ++i) {
EXPECT_EQ(expected_roots[i], forest->FindRoot(i));
}
}
Expand Down
Loading