Skip to content

Sharded return fix#2995

Open
wsmoses wants to merge 1 commit into
mainfrom
srf
Open

Sharded return fix#2995
wsmoses wants to merge 1 commit into
mainfrom
srf

Conversation

@wsmoses

@wsmoses wsmoses commented Jun 20, 2026

Copy link
Copy Markdown
Member

attempted fix for #2989

@wsmoses wsmoses requested a review from avik-pal June 20, 2026 22:59
Comment thread src/Tracing.jl
error("TODO(#2230): implement sharding")
ndev = Sharding.ndevices(sharding)
res = ConcretePJRTArray{T,N,ndev}(
ntuple(Returns(XLA.PJRT.AsyncEmptyBuffer), ndev), size(prev), Sharding.NoShardInfo()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
ntuple(Returns(XLA.PJRT.AsyncEmptyBuffer), ndev), size(prev), Sharding.NoShardInfo()
ntuple(Returns(XLA.PJRT.AsyncEmptyBuffer), ndev),
size(prev),
Sharding.NoShardInfo(),

mesh = Sharding.Mesh(reshape(addressable_devices, ndev), (:dev,))
sh = Sharding.DimsSharding(mesh, (1,), (:dev,))

x = Reactant.to_rarray(collect(Float64, 1:64); sharding = sh)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[JuliaFormatter] reported by reviewdog 🐶

Suggested change
x = Reactant.to_rarray(collect(Float64, 1:64); sharding = sh)
x = Reactant.to_rarray(collect(Float64, 1:64); sharding=sh)

Comment thread src/Tracing.jl
error("TODO(#2230): implement sharding")
ndev = Sharding.ndevices(sharding)
res = ConcretePJRTArray{T,N,ndev}(
ntuple(Returns(XLA.PJRT.AsyncEmptyBuffer), ndev), size(prev), Sharding.NoShardInfo()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is this really Sharding.NoShardInfo() here?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants