fix: relax device types and inputs for sharding (#930)

Author: avik-pal

src/Sharding.jl

@@ -315,7 +315,7 @@ function standardize_sharding(sharding::DimsSharding, x::Union{AbstractArray,Num
 end
 
 function (sharding::DimsSharding)(
-    client::XLA.AbstractClient, device::Nothing, x::Union{AbstractArray,Number}
+    client::XLA.AbstractClient, device, x::Union{AbstractArray,Number}
 )
     return (standardize_sharding(sharding, x))(client, device, x)
 end

src/Types.jl

@@ -289,7 +289,7 @@ function ConcreteIFRTArray(
                    arguments will be ignored."
         end
     end
-    sharded_data, sharding = sharding(client, device, data)
+    sharded_data, sharding = sharding(client, nothing, data)
     return ConcreteIFRTArray{T,N}(sharded_data, size(data), sharding)
 end