def coordinate_tensor(shape, axis):
  """Return a tensor with given shape containing coordinate along given axis.

  Args:
    shape: a Tensor representing the shape of the output Tensor
    axis: an integer

  Returns:
    A tensor with shape shape and type tf.int32, where each elements its
    coordinate along the given axis.
  """
  if axis < 0:
    axis = tf.size(shape) + axis  # Convert to positive for the one_hot indice

  r = tf.range(shape[axis])
  r_shape = tf.one_hot(
      axis, tf.size(shape), on_value=-1, off_value=1, dtype=tf.int32)
  return tf.zeros(shape, dtype=tf.int32) + tf.reshape(r, r_shape)
def coordinate_tensor ( shape , axis ) : if axis < 0 : axis = tf . size ( shape ) + axis r = tf . range ( shape [ axis ] ) r_shape = tf . one_hot ( axis , tf . size ( shape ) , on_value = - 1 , off_value = 1 , dtype = tf . int32 ) return tf . zeros ( shape , dtype = tf . int32 ) + tf . reshape ( r , r_shape )
Return a tensor with given shape containing coordinate along given axis.

  Args:
    shape: a Tensor representing the shape of the output Tensor
    axis: an integer

  Returns:
    A tensor with shape shape and type tf.int32, where each elements its
    coordinate along the given axis.