Searched refs:tpu_ordinal_function (Results 1 – 3 of 3) sorted by relevance
549 tpu_ordinal_function=None, argument596 if tpu_ordinal_function is None:597 tpu_ordinal_function = lambda index: -1 function604 tpu_ordinal=tpu_ordinal_function(index),624 tpu_ordinal_function=None): argument677 if tpu_ordinal_function is None:678 tpu_ordinal_function = self._default_ordinal_function689 if tpu_ordinal_function is None:690 tpu_ordinal_function = _ordinal_function_from_map723 tpu_ordinal=tpu_ordinal_function(index))
735 tpu_ordinal_function_impl = ctx.tpu_ordinal_function(host_id)768 per_host_sharded_inputs, tpu_ordinal_function=tpu_ordinal_function_impl)802 tpu_ordinal_function_impl = ctx.tpu_ordinal_function(host_id)837 tpu_ordinal_function=tpu_ordinal_function_impl))873 tpu_ordinal_function_impl = ctx.tpu_ordinal_function(host_id)927 tpu_ordinal_function=tpu_ordinal_function_impl)1023 tpu_ordinal_function=tpu_ordinal_function_impl,
565 def tpu_ordinal_function(self, host_id): member in _InternalTPUContext