# Path: mesh_tensorflow/utils.py
# def outside_all_rewrites():
#   def __init__(self, devices, init_usage=None):
#   def device_function(self, var):
# def create_host_call(model_dir):
#   def maybe_cast(tensor):
#   def host_call_fn(global_step, *args):
# def topology_rank(topology):
# def remove_summaries():
# class BalancedVariablePlacer(object):
# SCALAR_SUMMARIES_COLLECTION_KEY = 'mtf_scalar_summaries'


import collections
import copy
import functools
import itertools
import operator
import os
import re
import gin
import numpy as np
import six
import tensorflow.compat.v1 as tf
from mesh_tensorflow import utils
from six.moves import xrange  # pylint: disable=redefined-builtin
from tensorflow.python.ops.gen_nn_ops import conv3d_backprop_input_v2
from tensorflow.python.ops.nn_ops import conv3d_backprop_filter_v2
  def __init__(self, x, name=None):
    super(StopGradient, self).__init__(
        [x], x.mesh, name=name or "stop_gradient")
    self._outputs = [Tensor(self, x.shape, x.dtype)]

  def lower(self, lowering):
    lowering.set_tensor_lowering(self.outputs[0],
                                 lowering.tensors[self.inputs[0]])

  @property
  def has_gradient(self):
    return False


def stop_gradient(x):
  return StopGradient(x).outputs[0]


class ScalarSummaryOperation(Operation):
  """Similar to tf.Print."""

  def __init__(self, name, x):
    super(ScalarSummaryOperation, self).__init__(
        [x], x.mesh, name=name)
    if x.shape.dims:
      raise ValueError("ScalarSummaryOperation takes a scalar")
    self._outputs = [Tensor(self, x.shape, x.dtype)]

  def lower(self, lowering):
    lowered_input = lowering.tensors[self.inputs[0]].to_laid_out_tensor()

# Path: mesh_tensorflow/utils.py
# def outside_all_rewrites():
#   def __init__(self, devices, init_usage=None):
#   def device_function(self, var):
# def create_host_call(model_dir):
#   def maybe_cast(tensor):
#   def host_call_fn(global_step, *args):
# def topology_rank(topology):
# def remove_summaries():
# class BalancedVariablePlacer(object):
# SCALAR_SUMMARIES_COLLECTION_KEY = 'mtf_scalar_summaries'


# Path: mesh_tensorflow/ops.py
import collections
import copy
import functools
import itertools
import operator
import os
import re
import gin
import numpy as np
import six
import tensorflow.compat.v1 as tf
from mesh_tensorflow import utils
from six.moves import xrange  # pylint: disable=redefined-builtin
from tensorflow.python.ops.gen_nn_ops import conv3d_backprop_input_v2
from tensorflow.python.ops.nn_ops import conv3d_backprop_filter_v2

  def __init__(self, x, name=None):
    super(StopGradient, self).__init__(
        [x], x.mesh, name=name or "stop_gradient")
    self._outputs = [Tensor(self, x.shape, x.dtype)]

  def lower(self, lowering):
    lowering.set_tensor_lowering(self.outputs[0],
                                 lowering.tensors[self.inputs[0]])

  @property
  def has_gradient(self):
    return False


def stop_gradient(x):
  return StopGradient(x).outputs[0]


class ScalarSummaryOperation(Operation):
  """Similar to tf.Print."""

  def __init__(self, name, x):
    super(ScalarSummaryOperation, self).__init__(
        [x], x.mesh, name=name)
    if x.shape.dims:
      raise ValueError("ScalarSummaryOperation takes a scalar")
    self._outputs = [Tensor(self, x.shape, x.dtype)]

  def lower(self, lowering):
    lowered_input = lowering.tensors[self.inputs[0]].to_laid_out_tensor()

    tf.add_to_collection(utils.SCALAR_SUMMARIES_COLLECTION_KEY,