mirror of
https://github.com/ROCm/jax.git
synced 2025-04-25 12:06:10 +00:00
73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
# Copyright 2025 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.
|
|
"""`jax.experimental.transfer`: DCN cross slice transfer."""
|
|
|
|
import jax
|
|
from typing import Any, TYPE_CHECKING
|
|
from jax._src.lib import xla_client as _xc
|
|
from jax._src.util import use_cpp_class, use_cpp_method
|
|
|
|
class TransferConnection:
|
|
"""Represents a connection to exactly one peer."""
|
|
|
|
@use_cpp_method()
|
|
def _pull_flat(self, uuid, backend, xs_flat):
|
|
raise NotImplementedError()
|
|
|
|
def pull(self, uuid: int, xs: Any) -> Any:
|
|
"""Fetches a pytree of arrays from a remote device.
|
|
|
|
Args:
|
|
uuid: identifier for the request
|
|
xs: A pytree of ShapeDtypeStruct.
|
|
Returns:
|
|
A pytree of arrays.
|
|
"""
|
|
xs_flat, tree = jax.tree.flatten(xs)
|
|
if not xs_flat:
|
|
return xs
|
|
backend = next(iter(xs_flat[0].sharding.device_set)).client
|
|
return tree.unflatten(self._pull_flat(uuid, backend, xs_flat))
|
|
|
|
|
|
if not TYPE_CHECKING:
|
|
TransferConnection = use_cpp_class(_xc._xla.TransferConnection)(TransferConnection)
|
|
|
|
|
|
class TransferServer:
|
|
|
|
@use_cpp_method()
|
|
def address(self) -> str:
|
|
"""Returns the address that this server can be connected to with."""
|
|
raise NotImplementedError()
|
|
|
|
@use_cpp_method()
|
|
def _await_pull_flat(self, uuid, args: list[jax.Array]):
|
|
raise NotImplementedError()
|
|
|
|
@use_cpp_method()
|
|
def connect(self, address: str) -> TransferConnection:
|
|
"""Creates a connection to a remote server."""
|
|
raise NotImplementedError()
|
|
|
|
def await_pull(self, uuid: int, arrays: Any) -> Any:
|
|
"""Schedules a pytree of arrays to be fetched by a remote device."""
|
|
self._await_pull_flat(uuid, jax.tree.flatten(arrays)[0])
|
|
|
|
|
|
if not TYPE_CHECKING:
|
|
TransferServer = use_cpp_class(_xc._xla.TransferServer)(TransferServer)
|
|
|
|
start_transfer_server = _xc._xla.start_transfer_server
|