Home
last modified time | relevance | path

Searched refs:tpu_variables (Results 1 – 3 of 3) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/tests/
H A Dtpu_embedding_v1_checkpoint_test.py201 tpu_variables = checkpoint_utils.list_variables(
209 self.assertAllEqual(tpu_variables, cpu_variables)
H A Dtpu_embedding_v2_checkpoint_test.py359 tpu_variables = checkpoint_utils.list_variables(
367 self.assertAllEqual(tpu_variables, cpu_variables)
/aosp_15_r20/external/tensorflow/tensorflow/core/tpu/kernels/
H A Dtpu_functional_ops.cc1676 absl::flat_hash_map<uint64, Node*> tpu_variables; in ReplaceResourceArgsWithVarHandleOps() local
1696 if (enable_variable_deduplication && tpu_variables.contains(handle_fp) && in ReplaceResourceArgsWithVarHandleOps()
1698 Node* tpu_variable = tpu_variables.at(handle_fp); in ReplaceResourceArgsWithVarHandleOps()
1778 tpu_variables[handle_fp] = new_node; in ReplaceResourceArgsWithVarHandleOps()