@@ -75,7 +75,6 @@ function GPUClient(node_id=0, num_nodes=1, platform="gpu")
75
75
return Client (client)
76
76
end
77
77
78
-
79
78
function TPUClient (tpu_path:: String )
80
79
f = Libdl. dlsym (Reactant_jll. libReactantExtra_handle, " MakeTPUClient" )
81
80
refstr = Ref {Cstring} ()
@@ -103,18 +102,19 @@ function __init__()
103
102
104
103
@static if ! Sys. isapple ()
105
104
if isfile (" /usr/lib/libtpu.so" )
106
- dataset_dir = @get_scratch! (" libtpu" )
107
- if ! isfile (dataset_dir* " /libtpu.so" )
108
- Downloads. download (" https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20240829-py3-none-any.whl" , dataset_dir* " /tpu.zip" )
109
- run (` unzip -qq $(dataset_dir* " /tpu.zip" ) -d $(dataset_dir) /tmp` )
110
- run (` mv $(dataset_dir) /tmp/libtpu/libtpu.so $(dataset_dir) /libtpu.so` )
111
- rm (dataset_dir* " /tmp" , recursive= true )
112
- rm (dataset_dir* " /tpu.zip" , recursive= true )
113
- end
114
- try
115
- tpu = TPUClient (
116
- dataset_dir* " /libtpu.so"
105
+ dataset_dir = @get_scratch! (" libtpu" )
106
+ if ! isfile (dataset_dir * " /libtpu.so" )
107
+ Downloads. download (
108
+ " https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/wheels/libtpu-nightly/libtpu_nightly-0.1.dev20240829-py3-none-any.whl" ,
109
+ dataset_dir * " /tpu.zip" ,
117
110
)
111
+ run (` unzip -qq $(dataset_dir* " /tpu.zip" ) -d $(dataset_dir) /tmp` )
112
+ run (` mv $(dataset_dir) /tmp/libtpu/libtpu.so $(dataset_dir) /libtpu.so` )
113
+ rm (dataset_dir * " /tmp" ; recursive= true )
114
+ rm (dataset_dir * " /tpu.zip" ; recursive= true )
115
+ end
116
+ try
117
+ tpu = TPUClient (dataset_dir * " /libtpu.so" )
118
118
backends[" tpu" ] = tpu
119
119
default_backend[] = tpu
120
120
catch e
0 commit comments