Searched refs:get_tpu_context (Results 1 – 6 of 6) sorted by relevance
111 number_of_shards = tpu_function.get_tpu_context().number_of_shards164 num_replicas = tpu_function.get_tpu_context().number_of_shards
56 def get_tpu_context(): function
123 num_shards = tpu_function.get_tpu_context().number_of_shards
435 return tpu_function.get_tpu_context().number_of_shards is None
34 num_shards = tpu_function.get_tpu_context().number_of_shards
350 num_shards = tpu_function.get_tpu_context().number_of_shards