Skip to content

Commit 5752890

Browse files
LarryLansingGoogle-ML-Automation
authored andcommitted
Add shared ownership to tstring's VIEW type.
This change introduces `tstring::owner<T>`, a reference-counted wrapper for an object `T`. `tstring` can now be assigned a "shared view" using `assign_as_shared_view`, where the underlying data is owned by a `tstring::owner`. The `tstring` will increment the owner's reference count and decrement it when the `tstring` is deallocated or reassigned, ensuring the owner object remains alive as long as any `tstring` holds a shared view to its data. This is implemented by adding a `TStringOwnerCApi` pointer to the `TF_TString_View` struct and updating the relevant `TF_TString` functions to handle the reference counting. PiperOrigin-RevId: 832101409
1 parent 90ef1fa commit 5752890

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

jax_tpu_embedding/sparsecore/lib/core/input_preprocessing.cc

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -418,7 +418,7 @@ void PopulateOutput(TableState& state, PreprocessSparseDenseMatmulOutput& out,
418418
absl::Mutex& output_mutex) {
419419
state.stats_per_host.Flatten();
420420

421-
absl::MutexLock mutex_lock(output_mutex);
421+
absl::MutexLock mutex(output_mutex);
422422
out.lhs_row_pointers[state.stacked_table_name] =
423423
std::move(state.csr_arrays_per_host.row_pointers);
424424
out.lhs_embedding_ids[state.stacked_table_name] =

0 commit comments

Comments
 (0)