vec_tensordiag

vec_tensordiag(obj, axis=-1, extra_dims=1): 

LLM Docstring

Embed an array onto the diagonal of a higher-rank tensor.

extra_dims new axes (each the size of the first post-axis dimension) are prepended to the trailing block and the input is written along that diagonal via diag_indices, leaving zeros off-diagonal.

• obj: np.ndarray

  the array to place on the diagonal

• axis: int

  boundary between batch axes and the block to diagonalize

• extra_dims: int

  number of extra diagonal axes to add

• :returns: np.ndarray

  the diagonal-embedded tensor