Undoes the incorrect fix to the argument and instead fixes the type hints of the function for write_executable.call

PiperOrigin-RevId: 676055225

Author: lukebaumann

pathwaysutils/persistence/helper.py

@@ -146,7 +146,7 @@ def write_one_array(
       location, name, value, timeout.total_seconds()
   )
   write_executable = plugin_executable.PluginExecutable(write_request)
-  _, write_future = write_executable.call([[value]])
+  _, write_future = write_executable.call([value])
   return write_future
 
 

pathwaysutils/plugin_executable.py

@@ -15,7 +15,7 @@
 
 import concurrent.futures
 import threading
-from typing import List, Sequence, Tuple
+from typing import List, Sequence, Tuple, Union
 
 import jax
 from jax._src.interpreters import pxla
@@ -33,7 +33,7 @@ def __init__(self, prog_str: str):
 
   def call(
       self,
-      in_arr: Sequence[List[jax.Array]] = (),
+      in_arr: Sequence[Union[jax.Array, List[jax.Array]]] = (),
       out_shardings: Sequence[jax.sharding.XLACompatibleSharding] = (),
       out_avals: Sequence[jax.core.ShapedArray] = (),
       out_committed: bool = True,