mirror of
https://github.com/ROCm/jax.git
synced 2025-04-19 13:26:06 +00:00
b6361b3e76
Remove 2 redundant whitespaces mentioned in https://github.com/jax-ml/jax/pull/25056#pullrequestreview-2615387492. PiperOrigin-RevId: 727264168
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
# Copyright 2024 The JAX Authors.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
from contextlib import contextmanager
|
|
from jax._src import config
|
|
from jax._src.lib import xla_client
|
|
|
|
config_ext = xla_client._xla.config
|
|
|
|
|
|
@contextmanager
|
|
def extend_compute_type(c_type: str | None):
|
|
if c_type is None:
|
|
yield
|
|
return
|
|
|
|
prev = config.compute_on_context_manager.swap_local(c_type)
|
|
try:
|
|
if prev is not None and prev is not config_ext.unset and c_type != prev:
|
|
raise NotImplementedError(
|
|
'Nesting `compute_on` with different compute types is not supported'
|
|
f' yet. Current compute_on type: {prev}')
|
|
yield c_type
|
|
finally:
|
|
config.compute_on_context_manager.set_local(prev)
|
|
|
|
def current_compute_type() -> str | None:
|
|
return config.compute_on_context_manager.value
|
|
|
|
def _check_valid(c_type: str):
|
|
if (c_type not in {'device_host', 'device', 'tpu_sparsecore'}
|
|
and not c_type.startswith("gpu_stream:")):
|
|
raise ValueError(
|
|
f'Invalid compute type {c_type}. Current supported values '
|
|
'are `device_host`, `device`, `tpu_sparsecore`, and `gpu_stream:#`.')
|
|
|
|
@contextmanager
|
|
def compute_on(compute_type: str):
|
|
if not isinstance(compute_type, str):
|
|
raise TypeError("`compute_on`'s compute_type argument must be a string.")
|
|
_check_valid(compute_type)
|
|
|
|
with extend_compute_type(compute_type):
|
|
yield
|