__init__.py

# Copyright 2026 The TensorFlow MUSA Authors. All Rights Reserved.
#
# 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.
# ==============================================================================

"""TensorFlow MUSA plugin package.

This package provides:
- Automatic plugin loading on import
- Device discovery utilities for available MUSA devices
- Runtime configuration helpers for MUSA memory growth

Example usage:
    import tensorflow_musa as tf_musa

    # Plugin is automatically loaded on import
    devices = tf_musa.get_musa_devices()
"""

import logging

from ._graph_optimizer import (
    DISABLED_FUSION_PATTERNS_PARAM,
    MUSA_GRAPH_OPTIMIZER_NAME,
    clear_musa_disabled_fusion_patterns,
    clear_musa_graph_dump_config,
    disable_musa_graph_optimizer,
    disable_musa_fusion_patterns,
    disable_musa_graph_dump,
    enable_musa_graph_optimizer,
    enable_musa_graph_dump,
    get_musa_graph_dump_directory,
    get_musa_disabled_fusion_patterns,
    is_musa_graph_dump_enabled,
    is_musa_graph_dump_slim_enabled,
    is_musa_graph_dump_text_enabled,
    is_musa_graph_optimizer_enabled,
    set_musa_graph_dump_config,
    set_musa_disabled_fusion_patterns,
    set_musa_graph_optimizer_enabled,
)
from ._loader import get_musa_devices, get_musa_ops, is_plugin_loaded, load_plugin
from ._optimizers import apply_adam_mixed, apply_sparse_adam_mixed
from ._runtime_config import (
    disable_musa_telemetry,
    enable_musa_telemetry,
    get_musa_telemetry_health,
    is_musa_telemetry_enabled,
    set_musa_allow_growth,
    set_musa_telemetry_config,
)
from . import gradients, ops, raw_ops

# Package version
__version__ = "1.5.0"

# Load plugin automatically on import
_plugin_loaded = False

try:
    load_plugin()
    _plugin_loaded = True
except Exception as e:
    logging.warning(f"Failed to load MUSA plugin: {e}")
    logging.warning(
        "MUSA functionality will not be available. "
        "Please ensure the plugin is built and MUSA SDK is installed."
    )

# Public API
__all__ = [
    "__version__",
    "gradients",
    "ops",
    "raw_ops",
    "load_plugin",
    "get_musa_ops",
    "is_plugin_loaded",
    "get_musa_devices",
    "MUSA_GRAPH_OPTIMIZER_NAME",
    "DISABLED_FUSION_PATTERNS_PARAM",
    "set_musa_graph_optimizer_enabled",
    "enable_musa_graph_optimizer",
    "disable_musa_graph_optimizer",
    "is_musa_graph_optimizer_enabled",
    "set_musa_disabled_fusion_patterns",
    "disable_musa_fusion_patterns",
    "clear_musa_disabled_fusion_patterns",
    "get_musa_disabled_fusion_patterns",
    "set_musa_graph_dump_config",
    "enable_musa_graph_dump",
    "disable_musa_graph_dump",
    "clear_musa_graph_dump_config",
    "is_musa_graph_dump_enabled",
    "get_musa_graph_dump_directory",
    "is_musa_graph_dump_text_enabled",
    "is_musa_graph_dump_slim_enabled",
    "set_musa_allow_growth",
    "set_musa_telemetry_config",
    "enable_musa_telemetry",
    "disable_musa_telemetry",
    "is_musa_telemetry_enabled",
    "get_musa_telemetry_health",
    "apply_adam_mixed",
    "apply_sparse_adam_mixed",
    "MusaAdam",
]


def __getattr__(name):
    if name == "MusaAdam":
        from ._optimizers import _make_musa_adam_class

        cls = _make_musa_adam_class()
        globals()["MusaAdam"] = cls
        return cls
    raise AttributeError(f"module 'tensorflow_musa' has no attribute {name!r}")