internal

PiperOrigin-RevId: 728276097

.github/workflows/build.yml

@@ -172,12 +172,25 @@ jobs:
         sudo apt-get update
         sudo apt-get install -y protobuf-compiler
 
-        pip install .
+        pip install tensorflow
+
+        protoc -I=. --python_out=. $(find orbax/experimental/model/ -name "*.proto")
+
+        pip install -e .
         pip install .[testing] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
+
         if [[ "${{ matrix.jax-version }}" == "newest" ]]; then
           pip install -U jax jaxlib
         elif [[ "${{ matrix.jax-version }}" == "nightly" ]]; then
           pip install -U --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
         else
           pip install "jax==${{ matrix.jax-version }}" "jaxlib==${{ matrix.jax-version }}"
         fi
+    - name: Test with pytest
+      run: |
+        pytest orbax/experimental/model/core/python/*_test.py
+        echo "!!! done testing orbax/experimental/model/core/python/*_test.py \n"
+        pytest orbax/experimental/model/tf2obm/*_test.py
+        echo "!!! done testing orbax/experimental/model/tf2obm/*_test.py \n"
+        pytest orbax/experimental/model/jax2obm/*_test.py
+        echo "!!! done testing orbax/experimental/model/jax2obm/*_test.py \n"

model/orbax/experimental/model/conftest.py

@@ -0,0 +1,23 @@
+# Copyright 2024 The Orbax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""pytest config."""
+
+from typing import Any
+from absl import flags
+
+
+def pytest_configure(config: Any) -> None:
+  del config
+  flags.FLAGS.mark_as_parsed()

model/orbax/experimental/model/core/python/saved_model_proto/nodes_test.py

@@ -20,7 +20,7 @@
 from orbax.experimental.model.core.python.saved_model_proto import node_builder
 from orbax.experimental.model.core.python.saved_model_proto import nodes
 
-from .net.proto2.contrib.pyutil import compare
+from tensorflow.python.util.protobuf import compare
 from absl.testing import absltest
 
 Node = nodes.Node

model/orbax/experimental/model/jax2obm/jax_specific_info_test.py

@@ -19,7 +19,7 @@
 from orbax.experimental.model import core as obm
 from orbax.experimental.model.jax2obm import jax_specific_info
 from orbax.experimental.model.jax2obm import jax_supplemental_pb2
-from .net.proto2.contrib.pyutil import compare
+from tensorflow.python.util.protobuf import compare
 from google.protobuf import text_format
 from absl.testing import absltest
 

model/orbax/experimental/model/jax2obm/main_lib_test.py

@@ -29,7 +29,7 @@
 from orbax.experimental.model.jax2obm import jax_supplemental_pb2
 from orbax.experimental.model.jax2obm import main_lib
 
-from .net.proto2.contrib.pyutil import compare
+from tensorflow.python.util.protobuf import compare
 from google.protobuf import text_format
 from absl.testing import absltest
 

model/orbax/experimental/model/jax2obm/sharding_test.py

@@ -17,7 +17,7 @@
 import numpy as np
 from orbax.experimental.model import core as obm
 from orbax.experimental.model.jax2obm import sharding
-from .net.proto2.contrib.pyutil import compare
+from tensorflow.python.util.protobuf import compare
 from google.protobuf import text_format
 from absl.testing import absltest
 

model/orbax/experimental/model/tf2obm/tf_concrete_functions_to_obm_test.py

@@ -28,7 +28,7 @@
 from orbax.experimental.model.tf2obm.tf_concrete_functions_to_obm import TF_SAVED_MODEL_SUPPLEMENTAL_NAME
 import tensorflow as tf
 
-from .net.proto2.contrib.pyutil import compare
+from tensorflow.python.util.protobuf import compare
 from google.protobuf import text_format
 from absl.testing import absltest
 

model/pyproject.toml

@@ -46,4 +46,5 @@ all = [
 
 testing = [
     'pytest',
+    'flax',
 ]