Code update

PiperOrigin-RevId: 567448931
google-research · Sep 21, 2023 · 9bbddf9 · 9bbddf9
1 parent e72b535
commit 9bbddf9
Show file tree

Hide file tree

Showing 14 changed files with 91 additions and 410 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -23,10 +23,12 @@ keywords = [
 # Installed locally with `pip install -e .`
 dependencies = [
   "absl-py",
+  "chex",
   "clu",
+  "etils",
   "flax",
   "h5py",
-  # "grain@git+https://github.com/google/grain",
+  "grain@git+https://github.com/google/grain.git@b80f7066ce1f69317519bf64739e5ff9a463059a",
   "gin-config",
   "jax",
   "numpy",

diff --git a/swirl_dynamics/data/tfgrain_transforms.py b/swirl_dynamics/data/tfgrain_transforms.py
diff --git a/swirl_dynamics/data/tfgrain_transforms_test.py b/swirl_dynamics/data/tfgrain_transforms_test.py
diff --git a/swirl_dynamics/data/utils_test.py b/swirl_dynamics/data/utils_test.py
@@ -14,10 +14,13 @@
 
 import os
 
+from absl import flags
 from absl.testing import absltest
 from absl.testing import parameterized
 from swirl_dynamics.data import utils
 
+FLAGS = flags.FLAGS
+
 
 class UtilsTest(parameterized.TestCase):