Return the bucket id of the given tensor.

    Args:
      x (tf.Tensor): float32 of shape [length, depth]

    Returns:
      tf.Tensor: One-hot vector int64 of shape [heads, length, nb_buckets]
        containing the id of the bucket
423653
def get_gates(self, x):
    

    # The balance loss don't propagate to the rest of the network
    x = tf.stop_gradient(x)
    # [length, depth] * [depth, nb_vectors * replicat]
    x = tf.matmul(x, self.t_vectors)
    # [length, nb_vector * replicat]
    x = tf.sign(x)  # Get on which side of the hyperplane the keys are.

    # x = tf.reshape(x, [-1, nb_replicat, nb_vector])
    # [length, replicat, nb_vector] * [nb_vector, 2^nb_vector - 1]

    x = tf.matmul(x, self.t_group, transpose_b=True) / self.nb_hyperplanes
    # We get a similarity score for each of the group between [-1, 1]
    # [length, (replicat,) 2^nb_vector - 1]
    # Do an argmax to get the most likely group for each replicat
    x = tf.argmax(x, axis=-1)
    # [length(, replicat)]
    # One-hot for compatibility with the sparse dispatcher
    x = tf.one_hot(x, self.nb_buckets)
    # TODO(epot): Use a loss to force an even distribution
    return x
364053
def get_uid(self, idx):
        
        assert idx is not None

        if isinstance(idx, (int, float, str)):
            return self.uid[idx]

        ret = []
        for i in idx:
            tmp = self.uid.get(i, None)
            assert tmp is not None, (
                'Model <{}> does not have element <{}>'.format(self._name, i))
            ret.append(self.uid[i])

        return ret
960876
def xindex(self):
        
        try:
            return self._xindex
        except AttributeError:
            self._xindex = Index.define(self.x0, self.dx, self.shape[0])
            return self._xindex
1348651
def _generate_bucket_value(self, bucketing_id):
    

    ratio = float(self._generate_unsigned_hash_code_32_bit(bucketing_id)) / MAX_HASH_VALUE
    return math.floor(ratio * MAX_TRAFFIC_VALUE)
332454
def mtf_range(mesh, dim, dtype, name=None):
  
  dim = convert_to_dimension(dim)
  with tf.variable_scope(name, default_name="range"):
    if dtype == tf.bfloat16:
      # tf.range(dtype=bfloat16) gives the wrong shape.
      # TODO(noam): report the bug.
      tf_range = tf.cast(tf.range(dim.size), tf.bfloat16)
    else:
      tf_range = tf.range(dim.size, dtype=dtype)
    return import_tf_tensor(mesh, tf_range, shape=Shape([dim]))
834445
def get_range_vector(size: int, device: int) -> torch.Tensor:
    
    if device > -1:
        return torch.cuda.LongTensor(size, device=device).fill_(1).cumsum(0) - 1
    else:
        return torch.arange(0, size, dtype=torch.long)
1026901
def rank(tensor: BKTensor) -> int:
    """"""
    if isinstance(tensor, np.ndarray):
        return len(tensor.shape)

    return len(tensor[0].size())
910063
def forward_tensor(self, x):
        
        fwd = vec_to_tri(x, self.N)
        return tf.squeeze(fwd, axis=0) if self.squeeze else fwd
326428
public INDArray ndArrayFromTensor(TF_Tensor tensor) {
        int rank = TF_NumDims(tensor);

        int[] ndShape;
        if (rank == 0) {
            // scalar
            ndShape = new int[] { 1 };
        } else {
            ndShape = new int[rank];
            for (int i = 0; i < ndShape.length; i++) {
                ndShape[i] = (int) TF_Dim(tensor,i);
            }
        }

        int tfType = TF_TensorType(tensor);
        DataType nd4jType = typeFor(tfType);

        int length = ArrayUtil.prod(ndShape);
        INDArray array;
        if (nd4jType == DataType.UTF8) {
            String[] strings = new String[length];
            BytePointer data = new BytePointer(TF_TensorData(tensor)).capacity(TF_TensorByteSize(tensor));
            BytePointer str = new BytePointer((Pointer)null);
            SizeTPointer size = new SizeTPointer(1);
            TF_Status status = TF_NewStatus();
            for (int i = 0; i < length; i++) {
                long offset = data.position(8 * i).getLong();
                TF_StringDecode(data.position(8 * length + offset), data.capacity() - data.position(), str, size, status);
                if (TF_GetCode(status) != TF_OK) {
                    throw new IllegalStateException("ERROR: Unable to convert tensor " + TF_Message(status).getString());
                }
                strings[i] = str.position(0).capacity(size.get()).getString();
            }
            TF_DeleteStatus(status);
            array = Nd4j.create(strings);
        } else {
            Pointer pointer = TF_TensorData(tensor).capacity(length);
            Indexer indexer = indexerForType(nd4jType,pointer);
            DataBuffer d = Nd4j.createBuffer(indexer.pointer(),nd4jType,length,indexer);
            array = Nd4j.create(d,ndShape);
        }
        Nd4j.getAffinityManager().tagLocation(array, AffinityManager.Location.HOST);
        return array;
    }
469187
def get_variable(self, name, initializer):
    
    v = tf.get_variable(name, shape=initializer.shape,
                        initializer=(lambda shape, dtype, partition_info:
                                     initializer),
                        trainable=self.training)
    return v
268633
def np2model_tensor(a):
    ""
    dtype = model_type(a.dtype)
    res = as_tensor(a)
    if not dtype: return res
    return res.type(dtype)
651546
def MaxPooling(
        inputs,
        pool_size,
        strides=None,
        padding='valid',
        data_format='channels_last'):
    
    if strides is None:
        strides = pool_size
    layer = tf.layers.MaxPooling2D(pool_size, strides, padding=padding, data_format=data_format)
    ret = layer.apply(inputs, scope=tf.get_variable_scope())
    return tf.identity(ret, name='output')
1334957
def get_tensor_shape(self, tensor_name):
    
    tensor = self._name_to_tensor(tensor_name)

    if isinstance(tensor, mtf.Tensor):
      return tf.TensorShape(tensor.shape.to_integer_list)
    else:  # tf.Tensor
      return tensor.shape
514070
def GlobalAvgPooling(x, data_format='channels_last'):
    
    assert x.shape.ndims == 4
    data_format = get_data_format(data_format)
    axis = [1, 2] if data_format == 'channels_last' else [2, 3]
    return tf.reduce_mean(x, axis, name='output')
42428
def tensor_dimension_to_mesh_axis(self, tensor_dimension, mesh_shape):
    
    val = [i for i, mesh_dimension in enumerate(mesh_shape)
           if (tensor_dimension.name, mesh_dimension.name) in self._pairs]
    if len(val) > 1:
      raise ValueError(
          "Tensor dimension maps to multiple mesh dimensions"
          " tensor_dimension=%s mesh_shape=%s layout=%s"
          % (tensor_dimension, mesh_shape, self._pairs))
    return val[0] if val else None
169985
def import_tf_tensor(self, x, tf_x):
    
    return self.LaidOutTensor(self.make_slices(tf_x, x.shape))
1302554
def reconstruct(self, t:Tensor, x:Tensor=None):
        ""
        return self[0].reconstruct(t,x) if has_arg(self[0].reconstruct, 'x') else self[0].reconstruct(t)