(v, k=0)
return utils.tensor_to_ndarray(result)
Raises an error if input is not 1- or 2-d.
Raises an error if input is not 1- or 2-d.
Raises an error if input is not 1 - or 2 - d .
def diag(v, k=0):  # pylint: disable=missing-docstring
  """Raises an error if input is not 1- or 2-d."""
  v = asarray(v).data
  v_rank = tf.rank(v)

  v.shape.with_rank_at_most(2)

  # TODO(nareshmodi): Consider a utils.Assert version that will fail during
  # tracing time if the shape is known.
  tf.debugging.Assert(
      utils.logical_or(tf.equal(v_rank, 1), tf.equal(v_rank, 2)), [v_rank])

  def _diag(v, k):
    return utils.cond(
        tf.equal(tf.size(v), 0),
        lambda: tf.zeros([abs(k), abs(k)], dtype=v.dtype),
        lambda: tf.linalg.diag(v, k=k))

  def _diag_part(v, k):
    v_shape = tf.shape(v)
    v, k = utils.cond(
        utils.logical_or(
            utils.less_equal(k, -1 * utils.getitem(v_shape, 0)),
            utils.greater_equal(k, utils.getitem(v_shape, 1)),
        ), lambda: (tf.zeros([0, 0], dtype=v.dtype), 0), lambda: (v, k))
    result = tf.linalg.diag_part(v, k=k)
    return result

  result = utils.cond(
      tf.equal(v_rank, 1), lambda: _diag(v, k), lambda: _diag_part(v, k))
  return utils.tensor_to_ndarray(result)
def diag ( v , k = 0 ) : # pylint: disable=missing-docstring v = asarray ( v ) . data v_rank = tf . rank ( v ) v . shape . with_rank_at_most ( 2 ) # TODO(nareshmodi): Consider a utils.Assert version that will fail during # tracing time if the shape is known. tf . debugging . Assert ( utils . logical_or ( tf . equal ( v_rank , 1 ) , tf . equal ( v_rank , 2 ) ) , [ v_rank ] ) def _diag ( v , k ) : return utils . cond ( tf . equal ( tf . size ( v ) , 0 ) , lambda : tf . zeros ( [ abs ( k ) , abs ( k ) ] , dtype = v . dtype ) , lambda : tf . linalg . diag ( v , k = k ) ) def _diag_part ( v , k ) : v_shape = tf . shape ( v ) v , k = utils . cond ( utils . logical_or ( utils . less_equal ( k , - 1 * utils . getitem ( v_shape , 0 ) ) , utils . greater_equal ( k , utils . getitem ( v_shape , 1 ) ) , ) , lambda : ( tf . zeros ( [ 0 , 0 ] , dtype = v . dtype ) , 0 ) , lambda : ( v , k ) ) result = tf . linalg . diag_part ( v , k = k ) return result result = utils . cond ( tf . equal ( v_rank , 1 ) , lambda : _diag ( v , k ) , lambda : _diag_part ( v , k ) ) return utils . tensor_to_ndarray ( result )