mirror of
https://github.com/ROCm/jax.git
synced 2025-04-16 03:46:06 +00:00
72fbdb2eb5
The guarantee provided by this API is that the sharding of `x` and `y` will be the same! What the sharding will be is decided by GSPMD. The flow of sharding is bidirectional i.e. SPMD will choose what the sharding should be of `x` and `y` depending on it's propagation algorithm. It might end up being that the sharding chosen is not of `x` and `y` but something better. At the end of propagation `x` and `y` will be sharded alike. The API can be made variadic in the future i.e. `*args = shard_alike(*args)` depending on use cases. Fixes: https://github.com/google/jax/issues/15600 PiperOrigin-RevId: 592375936
18 lines
646 B
Python
18 lines
646 B
Python
# Copyright 2023 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 ific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from jax._src.shard_alike import (
|
|
shard_alike as shard_alike
|
|
)
|