__init__.py

# Copyright 2024 RecML authors <recommendations-ml@google.com>.
#
# 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_trainer import JaxState
from recml.core.training.jax_trainer import JaxTask
from recml.core.training.jax_trainer import JaxTrainer
from recml.core.training.jax_trainer import KerasState
from recml.core.training.keras_trainer import KerasTask
from recml.core.training.keras_trainer import KerasTrainer
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