"""Keras' DTensor library."""

_DTENSOR_API_ENABLED = False


# Conditional import the dtensor API, since it is currently broken in OSS.
if _DTENSOR_API_ENABLED:
  # pylint: disable=g-direct-tensorflow-import, g-import-not-at-top
  from tensorflow.dtensor import python as dtensor_api
else:
  # Leave it with a placeholder, so that the import line from other python file
  # will not break.
  dtensor_api = None



"""Keras' DTensor library."""

_DTENSOR_API_ENABLED = False


# Conditional import the dtensor API, since it is currently broken in OSS.
if _DTENSOR_API_ENABLED:
  from tensorflow.compat.v2.experimental import dtensor as dtensor_api  # pylint: disable=g-import-not-at-top
else:
  # Leave it with a placeholder, so that the import line from other python file
  # will not break.
  dtensor_api = None

Change keras to use dtensor public API.
Change keras to use dtensor public API.

PiperOrigin-RevId: 438605222

  """Keras' DTensor library."""
  
  _DTENSOR_API_ENABLED = False
  
  
  # Conditional import the dtensor API, since it is currently broken in OSS.
  if _DTENSOR_API_ENABLED:
+   from tensorflow.compat.v2.experimental import dtensor as dtensor_api  # pylint: disable=g-import-not-at-top
-   # pylint: disable=g-direct-tensorflow-import, g-import-not-at-top
-   from tensorflow.dtensor import python as dtensor_api
  else:
    # Leave it with a placeholder, so that the import line from other python file
    # will not break.
    dtensor_api = None
  
- 
- 
Change keras to use dtensor public API.
## Code Before:
"""Keras' DTensor library."""

_DTENSOR_API_ENABLED = False


# Conditional import the dtensor API, since it is currently broken in OSS.
if _DTENSOR_API_ENABLED:
  # pylint: disable=g-direct-tensorflow-import, g-import-not-at-top
  from tensorflow.dtensor import python as dtensor_api
else:
  # Leave it with a placeholder, so that the import line from other python file
  # will not break.
  dtensor_api = None



## Instruction:
Change keras to use dtensor public API.
## Code After:
"""Keras' DTensor library."""

_DTENSOR_API_ENABLED = False


# Conditional import the dtensor API, since it is currently broken in OSS.
if _DTENSOR_API_ENABLED:
  from tensorflow.compat.v2.experimental import dtensor as dtensor_api  # pylint: disable=g-import-not-at-top
else:
  # Leave it with a placeholder, so that the import line from other python file
  # will not break.
  dtensor_api = None
