Update copybara to follow the new code structure.

PiperOrigin-RevId: 745701568

Author: qinyiyan

mlrx/__init__.py

@@ -0,0 +1,43 @@
+# Copyright 2024 RecML authors <[email protected]>.
+#
+# 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.
+"""Public API for MLRX."""
+
+# pylint: disable=g-importing-member
+
+from recml.core import metrics
+from recml.core import utils
+from recml.core.data.iterator import DatasetIterator
+from recml.core.data.iterator import TFDatasetIterator
+from recml.core.data.preprocessing import PreprocessingMode
+from recml.core.data.tf_dataset_factory import DatasetShardingInfo
+from recml.core.data.tf_dataset_factory import TFDatasetFactory
+from recml.core.data.tf_dataset_factory import TFDSMetadata
+from recml.core.metrics.base_metrics import Metric
+from recml.core.training.core import Experiment
+from recml.core.training.core import run_experiment
+from recml.core.training.core import Trainer
+from recml.core.training.jax import JaxState
+from recml.core.training.jax import JaxTask
+from recml.core.training.jax import JaxTrainer
+from recml.core.training.jax import KerasState
+from recml.core.training.optax_factory import AdagradFactory
+from recml.core.training.optax_factory import AdamFactory
+from recml.core.training.optax_factory import OptimizerFactory
+from recml.core.training.partitioning import DataParallelPartitioner
+from recml.core.training.partitioning import ModelParallelPartitioner
+from recml.core.training.partitioning import NullPartitioner
+from recml.core.training.partitioning import Partitioner
+from recml.core.utils.types import Factory
+from recml.core.utils.types import FactoryProtocol
+from recml.core.utils.types import ObjectFactory

mlrx/metrics/__init__.py

@@ -26,7 +26,7 @@
 
 from absl import logging
 import jax
-from mlrx.utils import types
+from recml.core.utils import types
 import tensorflow.compat.v2 as tf
 import tensorflow_datasets as tfds
 

recml/core/__init__.py

@@ -0,0 +1,34 @@
+# Copyright 2024 RecML authors <[email protected]>.
+#
+# 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.
+"""Public API for metrics."""
+
+# pylint: disable=g-importing-member
+
+from recml.core.metrics.base_metrics import Metric
+from recml.core.metrics.base_metrics import scalar
+from recml.core.metrics.confusion_metrics import aucpr
+from recml.core.metrics.confusion_metrics import aucroc
+from recml.core.metrics.confusion_metrics import estimate_confusion_matrix
+from recml.core.metrics.confusion_metrics import f1_score
+from recml.core.metrics.confusion_metrics import fbeta_score
+from recml.core.metrics.confusion_metrics import precision
+from recml.core.metrics.confusion_metrics import precision_at_recall
+from recml.core.metrics.confusion_metrics import recall
+from recml.core.metrics.mean_metrics import accuracy
+from recml.core.metrics.mean_metrics import binary_accuracy
+from recml.core.metrics.mean_metrics import mean_squared_error
+from recml.core.metrics.mean_metrics import top_k_accuracy
+from recml.core.metrics.reduction_metrics import mean
+from recml.core.metrics.reduction_metrics import sum  # pylint: disable=redefined-builtin
+from recml.core.metrics.tools import MetricAccumulator

mlrx/data/iterator.py recml/core/data/iterator.pymlrx/data/iterator.py renamed to recml/core/data/iterator.py

@@ -15,7 +15,7 @@
 
 from absl.testing import absltest
 import jax.numpy as jnp
-from mlrx.metrics import base_metrics
+from recml.core.metrics import base_metrics
 
 
 class BaseMetricsTest(absltest.TestCase):

mlrx/data/preprocessing.py recml/core/data/preprocessing.pymlrx/data/preprocessing.py renamed to recml/core/data/preprocessing.py

@@ -20,7 +20,7 @@
 import jax.numpy as jnp
 import jaxtyping as jt
 import numpy as np
-from mlrx.metrics import base_metrics
+from recml.core.metrics import base_metrics
 
 # pylint: disable=redefined-outer-name
 EPSILON = 1e-7