Merge branch 'main' into brendt/ray

lanl · Oct 30, 2023 · 652151a · 652151a
2 parents 7229726 + aadb4ad
commit 652151a
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 4 deletions.
diff --git a/scico/flax/train/input_pipeline.py b/scico/flax/train/input_pipeline.py
@@ -26,7 +26,7 @@
 from .typed_dict import DataSetDict
 
 DType = Any
-KeyArray = Union[Array, jax.random.PRNGKeyArray]
+KeyArray = Union[Array, jax.Array]
 
 
 class IterateData:

diff --git a/scico/flax/train/state.py b/scico/flax/train/state.py
@@ -21,7 +21,7 @@
 from .typed_dict import ConfigDict, ModelVarDict
 
 ModuleDef = Any
-KeyArray = Union[Array, jax.random.PRNGKeyArray]
+KeyArray = Union[Array, jax.Array]
 PyTree = Any
 ArrayTree = optax.Params
 

diff --git a/scico/flax/train/steps.py b/scico/flax/train/steps.py
@@ -19,7 +19,7 @@
 from .state import TrainState
 from .typed_dict import DataSetDict, MetricsDict
 
-KeyArray = Union[Array, jax.random.PRNGKeyArray]
+KeyArray = Union[Array, jax.Array]
 PyTree = Any
 
 

diff --git a/scico/flax/train/trainer.py b/scico/flax/train/trainer.py
@@ -47,10 +47,11 @@
 from .typed_dict import ConfigDict, DataSetDict, MetricsDict, ModelVarDict
 
 ModuleDef = Any
-KeyArray = Union[Array, jax.random.PRNGKeyArray]
+KeyArray = Union[Array, jax.Array]
 PyTree = Any
 DType = Any
 
+
 # sync across replicas
 def sync_batch_stats(state: TrainState) -> TrainState:
     """Sync the batch statistics across replicas."""