Skip to content

Commit 0c10825

Browse files
jpuigcervercopybara-github
jpuigcerver
authored and
copybara-github
committed
Minor fix after jax changes.
PiperOrigin-RevId: 773571781
1 parent fef7f73 commit 0c10825

File tree

98 files changed

+104
-102
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

98 files changed

+104
-102
lines changed

‎vmoe/app.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

‎vmoe/checkpoints/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

‎vmoe/checkpoints/base.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

‎vmoe/checkpoints/base_test.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

‎vmoe/checkpoints/partitioned.py‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.
@@ -56,8 +56,10 @@
5656
IndexInfo = types.IndexInfo
5757
LazyArrayChunks = types.LazyArrayChunks
5858
MapResult = multiprocessing.pool.MapResult
59-
GSPMDSharding = jax.sharding.GSPMDSharding
6059
PyTree = Any
60+
Mesh = jax.sharding.Mesh
61+
NamedSharding = jax.sharding.NamedSharding
62+
PartitionSpec = jax.sharding.PartitionSpec
6163
Sharding = jax.sharding.Sharding
6264
Slice = types.Slice
6365
SliceNd = types.SliceNd
@@ -361,7 +363,7 @@ def _get_array_sharding_or_default(arr: jax.Array) -> Sharding:
361363
if hasattr(arr, 'sharding'):
362364
return arr.sharding
363365
else:
364-
return GSPMDSharding.get_replicated(jax.devices())
366+
return NamedSharding(Mesh(jax.devices(), 'x'), PartitionSpec())
365367

366368

367369
def _intersect_slicend(a: SliceNd, b: SliceNd) -> Optional[SliceNd]:

‎vmoe/checkpoints/partitioned_test.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

‎vmoe/checkpoints/serialization.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

‎vmoe/checkpoints/serialization_test.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

‎vmoe/checkpoints/types.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

‎vmoe/checkpoints/types_test.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2024 Google LLC.
1+
# Copyright 2025 Google LLC.
22
#
33
# Licensed under the Apache License, Version 2.0 (the "License");
44
# you may not use this file except in compliance with the License.

0 commit comments

Comments
 (0)