The Mesh TensorFlow dimensions associated with a tensor.

    Args:
      tensor_name: a string, name of a tensor in the graph.

    Returns:
      a [string], the names of Mesh TensorFlow dimensions.
420635
def get_tensor_mtf_dimension_names(self, tensor_name):
    
    tensor = self._name_to_tensor(tensor_name)
    if isinstance(tensor, mtf.Tensor):
      return tensor.shape.dimension_names
    else:  # tf.Tensor
      return []
114162
def dump_tf_graph(tfgraph, tfgraphdef):
    
    print("Nodes ({})".format(len(tfgraphdef.node)))
    f = "{:15} {:59} {:20} {}"
    print(f.format("kind", "scopeName", "shape", "inputs"))
    for node in tfgraphdef.node:
        scopename = node.name
        kind = node.op
        inputs = node.input
        shape = tf.graph_util.tensor_shape_from_node_def_name(tfgraph, scopename)
        print(f.format(kind, scopename, str(shape), inputs))
891290
def get_variable(self, name):
        
        name = get_op_tensor_name(name)[1]
        if len(self.vs_name):
            name_with_vs = self.vs_name + "/" + name
        else:
            name_with_vs = name
        return get_op_or_tensor_by_name(name_with_vs)
155102
def debug_tensor(tensor, name):
    
    logging.info(name)
    tensor = tensor.detach().float().cpu().numpy()
    logging.info(f'MIN: {tensor.min()} MAX: {tensor.max()} '
                 f'AVG: {tensor.mean()} STD: {tensor.std()} '
                 f'NAN: {np.isnan(tensor).sum()} INF: {np.isinf(tensor).sum()}')
936214
def convert_to_dimension(d):
  
  if d is None:
    return None
  if isinstance(d, Dimension):
    if not isinstance(d.name, str) or not isinstance(d.size, int):
      raise ValueError("Bad dimension %s" % (d,))
    return d
  name, size = d
  if isinstance(name, str) and isinstance(size, int):
    return Dimension(name, size)
  else:
    raise ValueError("could not convert %s to Dimension" % (d,))
946155
def _initialize_mesh_dimension_name_to_size(self, mesh_shape):
    
    mesh_dimension_name_to_size = {}  # {string: int}
    for mesh_dimension in mesh_shape.dims:
      mesh_dimension_name_to_size[mesh_dimension.name] = mesh_dimension.size
    return mesh_dimension_name_to_size
397687
def dimensionNames(self):
        
        nSubDims = len(self._subArrayShape)
        subArrayDims = ['SubDim{}'.format(dimNr) for dimNr in range(nSubDims)]
        return dimNamesFromDataset(self._h5Dataset) + subArrayDims
913808
def get_tensor_device(self, tensor_name):
    
    tensor = self._name_to_tensor(tensor_name)
    if isinstance(tensor, tf.Tensor):
      return tensor.device
    else:  # mtf.Tensor
      return None
1282444
def state_size(self):
    """"""
    hidden_size = tf.TensorShape(
        self._input_shape[:-1] + (self._output_channels,))
    return (hidden_size, hidden_size)
207576
def _initialize_variables(self):
    """"""
    # Initialize global variables.
    self._global_vars = {}  # Indexed by (MTF dimension, mesh dimension)
    for mtf_dimension_name in (
        self._layout_validator.splittable_mtf_dimension_names):
      for mesh_dimension_name in (
          self._layout_validator.mesh_dimension_name_to_size):
        name = _global_var_name(mtf_dimension_name, mesh_dimension_name)
        self._global_vars[(mtf_dimension_name, mesh_dimension_name)] = (
            self._model.NewBoolVar(name))

    # Initialize local variables.
    self._local_vars = {}  # Indexed by (tensorflow dimension set), then name of
    # assignment.
    for mtf_dimension_set in self._mtf_dimension_sets:
      self._local_vars[mtf_dimension_set] = {}
      for assignment in self._assignments[mtf_dimension_set]:
        # TODO(joshuawang): Avoid hash collision no matter what dimension names
        # are; don't hash by this local var name, swap to using a tuple encoding
        # of the full assignment instead.
        name = _local_var_name(mtf_dimension_set, assignment)
        self._local_vars[mtf_dimension_set][name] = (
            self._model.NewBoolVar(name))

    # Initialize memory variable. We need a crude upper bound on memory, so we
    # use the total size of all tensors under the empty assignment.
    # NOTE(joshuawang): This bound could be improved by factoring in the
    # schedule.
    memory_upper_bound = 0
    for tensor_name in self._graph.get_all_tensor_names():
      if self._graph.is_tensor_on_canonical_device(tensor_name):
        memory_upper_bound += int(self._graph.get_tensor_size(tensor_name))
    self._memory_var = self._model.NewIntVar(0, memory_upper_bound, "z")
972981
def get_tensor_info(self):
    """"""
    # Add the additional length dimension to every shape

    def add_length_dim(tensor_info):
      return feature_lib.TensorInfo(
          shape=(self._length,) + tensor_info.shape,
          dtype=tensor_info.dtype,
      )

    tensor_info = super(SequenceDict, self).get_tensor_info()
    return utils.map_nested(add_length_dim, tensor_info)
1248435
def get_operation_output_names(self, operation_name):
    
    for output_tensor in self._name_to_operation(operation_name).outputs:
      yield output_tensor.name
591859
def _infer_shape(self, dimensions):
    
    # Size of input
    n = np.prod(dimensions)
    # Size of output where defined
    m = np.prod(abs(np.array(self._shape)))
    # Replace wildcard
    v = np.array(self._shape)
    v[v == -1] = n // m
    return tuple(v)
581017
def get_op_or_tensor_by_name(name):
    
    G = tfv1.get_default_graph()

    def f(n):
        if len(n) >= 3 and n[-2] == ':':
            return G.get_tensor_by_name(n)
        else:
            return G.get_operation_by_name(n)

    if not isinstance(name, list):
        return f(name)
    else:
        return list(map(f, name))
1016840
def model_sizes(m:nn.Module, size:tuple=(64,64))->Tuple[Sizes,Tensor,Hooks]:
    ""
    with hook_outputs(m) as hooks:
        x = dummy_eval(m, size)
        return [o.stored.shape for o in hooks]
957481
def get_tensor_info(self):
    """"""
    return {
        feature_key: feature.get_tensor_info()
        for feature_key, feature in self._feature_dict.items()
    }
1038126
def metric_tensor(self) -> np.ndarray:
        
        return dot(self._matrix, self._matrix.T)