Home
last modified time | relevance | path

Searched refs:_get_replica_numpy (Results 1 – 5 of 5) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/python/tpu/tests/
H A Dtpu_embedding_v2_sequence_feature_test.py56 self._get_replica_numpy(output[0], strategy, 0).shape, (2, 3, 4))
58 self._get_replica_numpy(output[1], strategy, 0).shape, (2, 3, 4))
60 self._get_replica_numpy(output[2], strategy, 0).shape, (2, 3, 2))
99 self._get_replica_numpy(output[0], strategy, 0).shape, (2, 3, 4))
101 self._get_replica_numpy(output[1], strategy, 0).shape, (2, 3, 4))
103 self._get_replica_numpy(output[2], strategy, 0).shape, (2, 3, 2))
H A Dtpu_embedding_v2_valid_input_test.py70 first = self._get_replica_numpy(
72 second = self._get_replica_numpy(embedding_only(data), strategy, 0)
162 activations0 = self._get_replica_numpy(activations, strategy, 0)
163 per_device_activations0 = self._get_replica_numpy(
205 no_weights0 = self._get_replica_numpy(no_weights_activations, strategy, 0)
206 weights0 = self._get_replica_numpy(weights_activations, strategy, 0)
H A Dtpu_embedding_v2_enqueue_mode_test.py70 activations_oc0 = self._get_replica_numpy(activations_oc, strategy, 0)
71 activations0 = self._get_replica_numpy(activations, strategy, 0)
175 before_gradient_apply0 = self._get_replica_numpy(before_gradient_apply,
177 after_gradient_apply0 = self._get_replica_numpy(after_gradient_apply,
H A Dtpu_embedding_v2_correctness_sequence_feature_test.py85 before_update = self._get_replica_numpy(
87 after_update = self._get_replica_numpy(embedding_only(data), strategy, 0)
H A Dtpu_embedding_base_test.py578 def _get_replica_numpy(self, structured, strategy, replica_id): member in TPUEmbeddingBaseTest