Format Julia code (#1477)

Co-authored-by: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>

Author: github-actions[bot]

src/ConcreteRArray.jl

@@ -371,12 +371,15 @@ function Base.setindex!(a::ConcreteIFRTArray, v, args::Vararg{Int,N}) where {N}
     return a
 end
 
-@inline function Base.similar(::Type{<:ConcretePJRTArray}, ::Type{S}, dims::Dims;
-                      client::Union{Nothing,XLA.PJRT.Client}=nothing,
-                      idx::Union{Int,Nothing}=nothing,
-                      device::Union{Nothing,XLA.PJRT.Device}=nothing,
-                      sharding::Sharding.AbstractSharding=Sharding.NoSharding()
-    ) where {S}
+@inline function Base.similar(
+    ::Type{<:ConcretePJRTArray},
+    ::Type{S},
+    dims::Dims;
+    client::Union{Nothing,XLA.PJRT.Client}=nothing,
+    idx::Union{Int,Nothing}=nothing,
+    device::Union{Nothing,XLA.PJRT.Device}=nothing,
+    sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
+) where {S}
     client = client === nothing ? XLA.default_backend() : client
 
     if idx isa Int && device === nothing
@@ -385,7 +388,9 @@ end
 
     sdata, sharding = sharding(client, device, S, dims)
 
-    return ConcretePJRTArray{S,length(dims),length(sdata),typeof(sharding)}(sdata, dims, sharding)
+    return ConcretePJRTArray{S,length(dims),length(sdata),typeof(sharding)}(
+        sdata, dims, sharding
+    )
 end
 
 function Base.similar(

src/Sharding.jl

@@ -418,9 +418,7 @@ function (sharding::NamedSharding)(
     return data, ShardInfo(sharding, device_to_array_slices)
 end
 
-function (sharding::NamedSharding)(
-    client::XLA.PJRT.Client, _, S::Type, dims::Dims
-)
+function (sharding::NamedSharding)(client::XLA.PJRT.Client, _, S::Type, dims::Dims)
     if !issorted(sharding.mesh.logical_device_ids)
         error("PJRT doesn't support non-iota meshes. Use IFRT instead.")
     end
@@ -431,7 +429,13 @@ function (sharding::NamedSharding)(
 
     data = ntuple(length(sharding.mesh)) do i
         Base.@_inline_meta
-        Base.similar(XLA.PJRT.AsyncBuffer, S, Dims(length.(device_to_array_slices[i])); client, device=XLA.get_device(client, sharding.mesh.device_ids[i]))
+        Base.similar(
+            XLA.PJRT.AsyncBuffer,
+            S,
+            Dims(length.(device_to_array_slices[i]));
+            client,
+            device=XLA.get_device(client, sharding.mesh.device_ids[i]),
+        )
     end
 
     return data, ShardInfo(sharding, device_to_array_slices)
@@ -774,7 +778,6 @@ function (sharding::Replicated)(client::XLA.PJRT.Client, dev, S::Type, dims::Dim
     return (NamedSharding(sharding, length(dims)))(client, dev, S, dims)
 end
 
-
 function sharding_to_array_slices(sharding::Replicated, size_x; kwargs...)
     return sharding_to_array_slices(
         NamedSharding(sharding, length(size_x)), size_x; kwargs...
@@ -963,13 +966,17 @@ function (sharding::HloSharding)(
     return data, ShardInfo(sharding, device_to_array_slices)
 end
 
-function (sharding::HloSharding)(
-    client::XLA.PJRT.Client, ::Nothing, S::Type, dims::Dims
-)
+function (sharding::HloSharding)(client::XLA.PJRT.Client, ::Nothing, S::Type, dims::Dims)
     device_to_array_slices = sharding_to_array_slices(sharding, dims; client)
 
-    data = ntuple(length(sharding.mesh)) do i    
-        Base.similar(XLA.PJRT.AsyncBuffer, S, Dims(length.(device_to_array_slices[i])); client, device=XLA.get_device(client, sharding.mesh.device_ids[i]))
+    data = ntuple(length(sharding.mesh)) do i
+        Base.similar(
+            XLA.PJRT.AsyncBuffer,
+            S,
+            Dims(length.(device_to_array_slices[i]));
+            client,
+            device=XLA.get_device(client, sharding.mesh.device_ids[i]),
+        )
     end
 
     return data, ShardInfo(sharding, device_to_array_slices)

src/xla/PJRT/AsyncBuffer.jl

@@ -18,4 +18,4 @@ end
 
 @inline function Base.similar(::Type{AsyncBuffer}, args...; kwargs...)
     return AsyncBuffer(Base.similar(Buffer, args...; kwargs...)::Buffer, nothing)
-end
+end

src/xla/PJRT/Buffer.jl

@@ -73,11 +73,14 @@ function Base.similar(a::Buffer, dims::Dims)
     return Buffer(buffer)
 end
 
-@inline function Base.similar(::Type{Buffer}, S::Type, dims::Dims;
-                      client::Union{Nothing,XLA.PJRT.Client}=nothing,
-                      idx::Union{Int,Nothing}=nothing,
-                      device::Union{Nothing,XLA.PJRT.Device}=nothing,
-                      )
+@inline function Base.similar(
+    ::Type{Buffer},
+    S::Type,
+    dims::Dims;
+    client::Union{Nothing,XLA.PJRT.Client}=nothing,
+    idx::Union{Int,Nothing}=nothing,
+    device::Union{Nothing,XLA.PJRT.Device}=nothing,
+)
     client = client === nothing ? XLA.default_backend() : client
 
     if device === nothing
@@ -108,7 +111,7 @@ end
 end
 
 function Base.similar(a::Buffer, S::Type, dims::Dims)
-    Base.similar(Buffer, S, dims; client=XLA.client(a), device=XLA.device(a))
+    return Base.similar(Buffer, S, dims; client=XLA.client(a), device=XLA.device(a))
 end
 
 @inline function free_buffer(buffer::Buffer)