cr.nimble.set_diagonal

cr.nimble.set_diagonal(A, value)[source]

Sets the diagonal elements to a specific value

Parameters

  • A (jax.numpy.ndarray) – A 2D matrix

  • value (float) – A value to be added to the diagonal elements

Returns

Matrix with updated diagonal

Return type

(jax.numpy.ndarray)